BOJ 2814 : 최소인수

문제

가장 작은 소인수가 $P$인 양의 정수 중에서 $N$번째로 작은 수를 구하는 프로그램을 작성하시오.

입력

첫째 줄에 $N$과 $P$가 주어진다. $(1 \le N, P \le 10^9)$

$P$는 항상 소수이다.

출력

첫째 줄에 가장 작은 소인수가 $P$인 양의 정수 중 $N$번째로 작은 수를 출력한다.

만약, 그러한 수가 $10^9$를 넘을 경우에는 0을 출력한다.

알고리즘 분류

• 수학
• 정수론
• 이분 탐색
• 많은 조건 분기
• 소수 판정
• 에라토스테네스의 체
• 포함 배제의 원리

풀이

일단 $P$의 범위가 너무 크니($1 \le P \le 10^9$) 줄이고 시작해보자.

한번 $P \gt \sqrt{10^9}$인 경우를 생각해보자. 우선 $N = 1$이면 정답은 $P$이다. $N = 2$인 경우는, 가장 작은 소인수가 $P$인 양의 정수 중 2번째로 작은 수는 $P^2$이고 $P^2 \gt 10^9$이므로 정답은 $0$이다. 그리고 $N > 2$일 때도 마찬가지로 정답은 $0$이다.

결국 $P \le \sqrt{10^9} (\approx 31623)$인 경우만 잘 해결하면 된다. 시간복잡도를 상당히 줄일 수 있다.

이 상태에서 문제를 직접적으로 풀어내기는 어려우니, 최적화 문제를 결정 문제로 변환하고 이분 탐색(파라메트릭 서치)으로 풀어보자.

$f(x) = (x$ 이하의 양의 정수 중 가장 작은 소인수가 $P$인 수의 개수$)$ 라고 정의해보자. 이때 $f(x') = N$이고 $f(x' - 1) = N - 1$인 양의 정수 $x'$를 찾으면 그게 정답이 된다. $f(x)$는 단조 증가 함수이므로 이분 탐색으로 풀이가 가능하며, 본래의 문제를 직접적으로 푸는 것보다 쉽게 풀 수 있다.

그럼 이제 $f(x)$의 값을 계산하는 방법을 알아내보자. $f(x)$의 값을 바로 알아내기는 어렵고 문제를 쪼개서 생각하면 된다. 문제를 쪼개보면 $($가장 작은 소인수가 $P$인 수의 개수$) = ($소인수 $P$를 가진 수의 개수$) - ($소인수 $P$를 가지며 $P$보다 작은 소인수도 가진 수의 개수$)$ 이다. 이렇게 각각을 구해서 빼주면 $f(x)$의 값을 구할 수 있다.

우선 소인수 $P$를 가진 수의 개수는 $\lfloor x / P \rfloor$로 간단하게 계산할 수 있다. 소인수 $P$를 가지며 $P$보다 작은 소인수도 가진 수의 개수만 잘 구해주면 된다.

이를 구하기 위해서 일단 $P$ 미만의 소수로 $2, 3, \dots , q$가 존재한다고 해보자. (실제 구현할 때는 에라토스테네스의 체로 구하면 된다)

그러면 소인수 $P$를 가진 $x$ 이하의 양의 정수들 중에서, $|(P$보다 작은 소인수를 가진 수의 집합$)| = |($소인수 2를 가진 수의 집합$) \cup ($소인수 3을 가진 수의 집합$) \cup \dots \cup ($소인수 $q$를 가진 수의 집합$)|$ 으로 구할 수 있다.

일단 합집합의 크기는 포함 배제의 원리를 이용하면 구할 수 있다. 그리고 포함 배제의 원리를 적용할 때, 특정 소인수들을 가진 수의 집합의 크기를 구할 수 있어야 하는데 그건 간단하다. 소인수 $P$를 가진 $x$ 이하의 양의 정수들 중에서, 소인수 $2, 3, \dots , r$도 가지고 있는 수의 개수는 $\lfloor x / (2 \times 3 \times \dots \times r \times P) \rfloor$로 간단하게 계산된다.

실제 구현할 때는 재귀 호출 형태로 소인수 집합의 경우의 수를 탐색하면 된다. 다만 필요 없는 경우를 스킵해서 탐색 속도를 빠르게 만들어야 통과할 수 있다. 소인수들을 선택하는 재귀 호출 과정에서, 현재까지 선택된 소인수들의 곱이 만약 $10^9$를 초과한다면 정답의 범위를 넘어버린다. 여기에 또 다른 소인수를 추가하여 탐색할 필요가 없다.

코드

먼저 직접 구현을 시도해보고, 코드를 참고하는걸 권장합니다.

import math

prime_limit = math.ceil(math.sqrt(10 ** 9))
is_prime = [True] * (prime_limit + 1)
(is_prime[1], is_prime[2]) = (False, True)

i = 2

while i * i <= prime_limit:
    if is_prime[i] == False:
        i += 1
        continue

    j = i * i

    while j <= prime_limit:
        is_prime[j] = False
        j += i

    i += 1

(n, p) = map(int, input().split())

if n == 1:
    print(p)
    exit(0)

if p * p > 10 ** 9:
    print(0)
    exit(0)

primes = []

for i in range(2, prime_limit + 1):
    if is_prime[i] and i < p:
        primes.append(i)

def solve(remain_product_count, primes_start_index, now_val, limit_val):
    ans = 0

    for i in range(primes_start_index, len(primes)):
        if remain_product_count > 1:
            next_val = now_val * primes[i]

            if next_val <= limit_val:
                temp_ans = solve(remain_product_count - 1, i + 1, next_val, limit_val)
                ans += temp_ans

                if temp_ans == 0:
                    break
            else:
                break
        else:
            final_val = now_val * primes[i]

            if final_val <= limit_val:
                ans += limit_val // final_val
            else:
                break

    return ans

(left, right) = (1, 10 ** 9 + 1)

while left < right:
    mid = (left + right) // 2
    has_p_count = mid // p
    has_p_and_lower_prime_count = 0

    for choice in range(1, 10):
        factor = (1 if choice % 2 == 1 else -1)
        has_p_and_lower_prime_count += solve(choice, 0, p, mid) * factor
    
    has_p_as_is_lowest_count = has_p_count - has_p_and_lower_prime_count

    if has_p_as_is_lowest_count < n:
        left = mid + 1
    elif has_p_as_is_lowest_count > n:
        right = mid - 1
    else:
        prime_check = True

        for prime in primes:
            if mid % prime == 0:
                prime_check = False
                break

        if mid % p != 0:
            prime_check = False

        if prime_check:
            left = right = mid
        else:
            right = mid - 1

print(left if left <= 10 ** 9 else 0)