[PS] BOJ 22943 / 수

[PS] BOJ 22943 / 수
Thumbnail: Photo by Heamn Aulakh (Unsplash)
문제 링크: https://www.acmicpc.net/problem/22943

범위 안의 소수를 미리 구한 뒤, 백트래킹으로 가능한 경우의 수를 세면 됩니다.

풀이

1) 소수 찾기

​주어진 범위 안의 소수를 찾기 위해서는, 에라토스테네스의 체를 사용하면 됩니다.

문제에서는 0~9까지의 서로 다른 $K$개의 수를 골라 사용하므로, $10^K$보다 작은 소수를 구하면 됩니다.

SIZE = 10 ** K

# 1. 에라토스테네스의 체
sieve = [True for _ in range(SIZE)]
sieve[0] = sieve[1] = False
for i in range(2, ceil(sqrt(SIZE))):
    if sieve[i]:
        for j in range(i * 2, SIZE, i):
            sieve[j] = False
# 계산한 소수를 배열에 따로 저장합니다.
primes = []
for i in range(2, SIZE):
    if sieve[i]:
        primes.append(i)

에라토스테네스의 체

  • sieve 배열은 에라토스테네스의 체로 사용하는 $1 \cdots 10^k - 1$의 수의 배열입니다. 범위 안의 어떤 수 $x$가 소수라면 True, 그렇지 않으면 False가 저장됩니다.
  • primes배열은 앞서 구한 소수를 저장한 배열입니다.

2) 조건에 맞는 수 판별하기

문제에서는 0~9까지의 서로 다른 $K$개의 수를 골라 만든 수 $x$가 다음의 두 조건을 모두 만족해야 합니다.

  1. $x$를 서로 다른 두 개의 소수의 합으로 나타낼 수 있는 경우
  2.  $x$를 $M$으로 나누어 떨어지지 않을때까지 나눈 수가 두 개의 소수의 곱인 경우. 이 때, 두 개의 소수가 같아도 된다.

각각을 판단하기 위해서, 다음을 미리 계산해 둬야 합니다.

  • $10^K$보다 작은 서로 다른 두 개의 소수의 합
  • $10^K$보다 작은 두 소수의 곱 (같은 소수끼리의 곱도 가능)
# 2. 어떤 서로 다른 두 소수의 합과 곱으로 만들 수 있는 수를 범위 안에서 미리 다 만들어 두기
sum_of_primes = [False for _ in range(SIZE)]
mul_of_primes = [False for _ in range(SIZE)]
for i in range(len(primes)):
    for j in range(len(primes)):
        if i != j:
            value = primes[i] + primes[j]
            if value < SIZE:
                sum_of_primes[value] = True
        
        value = primes[i] * primes[j]
        if value < SIZE:
            mul_of_primes[value] = True
        elif primes[i] + primes[j] >= SIZE:
            break

두 소수의 합과 곱 계산해두기

이후, 계산해 둔 정보를 가지고 어떤 수 $x$가 조건에 맞는지 판단하기 위한 함수 check를 정의합니다.

# 2. 각 수에 대해 조건에 맞는지 검사하는 로직
def check(x):
    # 1) 서로 다른 두 개의 소수의 합으로 나타낼 수 있는 경우
    if not sum_of_primes[x]:
        return False
    
    # 2) M으로 나누어 떨어지지 않을때까지 나눈 수가 두 개의 소수의 곱인 경우, 이 때, 두 개의 소수가 같아도 된다.
    while x % M == 0:
        x //= M
    if not mul_of_primes[x]: # 두 개의 소수의 곱 = 합성수인 경우
        return False

    return True

어떤 수 x가 조건에 맞는지 검사하는 함수

3) 조건에 맞는 수의 개수 세기

이제, 0~9까지의 서로 다른 $K$개의 수를 골라 만든 수 $x$ 중 위 2가지 조건을 만족하는 수의 개수를 세야 합니다. 백트래킹을 통해 가능한 모든 수의 조합을 만들어주면 됩니다.

# 3. 0 ~ 9 까지의 수를 K개 골라 만들 수 있는 모든 수 중 조건에 맞는 수를 추린다.
digits = [0, 0, 0, 0, 0]
def make_number():
    n = 0
    base = 1
    for i in range(K - 1, -1, -1):
        n += digits[i] * base
        base *= 10
    return n

visit = [False for _ in range(10)]
def backtrack(digit):
    res = 0
    if digit == K:
        return 1 if check(make_number()) else 0
    for i in range(10):
        if visit[i] or (digit == 0 and i == 0):
            continue
        digits[digit] = i
        visit[i] = True
        res += backtrack(digit + 1)
        visit[i] = False
    return res
     
print(backtrack(0))

백트래킹으로 조건에 맞는 수의 개수 세기

전체 코드

from math import ceil, sqrt
input = open(0).readline
K, M = map(int, input().split())
SIZE = 10 ** K

# 1. 에라토스테네스의 체
sieve = [True for _ in range(SIZE)]
sieve[0] = sieve[1] = False
primes = []
for i in range(2, ceil(sqrt(SIZE))):
    if sieve[i]:
        for j in range(i * 2, SIZE, i):
            sieve[j] = False
for i in range(2, SIZE):
    if sieve[i]:
        primes.append(i)

# 2. 어떤 서로 다른 두 소수의 합과 곱으로 만들 수 있는 수를 범위 안에서 미리 다 만들어 두기
sum_of_primes = [False for _ in range(SIZE)]
mul_of_primes = [False for _ in range(SIZE)]
for i in range(len(primes)):
    for j in range(len(primes)):
        if i != j:
            value = primes[i] + primes[j]
            if value < SIZE:
                sum_of_primes[value] = True
        
        value = primes[i] * primes[j]
        if value < SIZE:
            mul_of_primes[value] = True
        elif primes[i] + primes[j] >= SIZE:
            break

# 2. 각 수에 대해 조건에 맞는지 검사하는 로직
def check(x):
    # 1) 서로 다른 두 개의 소수의 합으로 나타낼 수 있는 경우
    if not sum_of_primes[x]:
        return False
    
    # 2) M으로 나누어 떨어지지 않을때까지 나눈 수가 두 개의 소수의 곱인 경우, 이 때, 두 개의 소수가 같아도 된다.
    while x % M == 0:
        x //= M
    if not mul_of_primes[x]: # 두 개의 소수의 곱 = 합성수인 경우
        return False

    return True

# 3. 0 ~ 9 까지의 수를 K개 골라 만들 수 있는 모든 수 중 조건에 맞는 수를 추린다.
digits = [0, 0, 0, 0, 0]
def make_number():
    n = 0
    base = 1
    for i in range(K - 1, -1, -1):
        n += digits[i] * base
        base *= 10
    return n

visit = [False for _ in range(10)]
def backtrack(digit):
    res = 0
    if digit == K:
        return 1 if check(make_number()) else 0
    for i in range(10):
        if visit[i] or (digit == 0 and i == 0):
            continue
        digits[digit] = i
        visit[i] = True
        res += backtrack(digit + 1)
        visit[i] = False
    return res
     
print(backtrack(0))

solution.py