[BOJ #15965] K번째 소수

BOJ #15965, K번째 소수 문제입니다.

 

문제는 간단합니다. 입력으로 K가 주어지면 K번째 소수를 찾아 출력하면 됩니다. 그러나 K의 범위에 따라 서브태스크가 여럿 존재하는 것으로 보아 K의 범위가 문제풀이에 중요하게 작용할 것이라고 예측할 수 있습니다. 이 문제에서는 최대 50만번째 소수를 2초 안에 찾아내야 하므로, 일반적인 방법으로는 시간이 부족할 것으로 보입니다.

 

1. 소수 구하기

1번째 방법

2 이상의 자연수 N에 대해, N이 1과 N을 제외하고 어떤 자연수로도 나누어 떨어지지 않을 때 그 수를 소수라고 합니다. 가장 원시적인 방법으로는 다음 코드가 가능하겠습니다.

def is_prime(N):
    if N == 2:
        return True
    else:
        for i in range(2, N):
            if not N % i:
                return False
        return True

위 코드를 보면, N이 소수인지 확인하기 위해서 대략 N번의 연산이 필요함을 알 수 있습니다. 미리 말씀드리자면, 50만번째 소수는 7368787입니다. 해당 숫자가 소수임을 검증하는 데에만 해도 대략 730만 번의 연산이 필요합니다. 등차수열의 합 공식을 이용해 1부터 7368787까지의 자연수들을 모두 돌릴 때 연산 횟수를 추정해 보면, 대략 7368787 * 7368787 / 2 = 27149510925684(약 271495초)가 됩니다. 절대 시간 내에 풀 수 없습니다.

 

2번째 방법

소수에 대해서 살짝 알고 계시다면, 2를 제외하면 모든 소수는 홀수이므로 위 코드를 다음과 같은 방식으로 개선할 수 있습니다.

if N == 2:
        return True
    else:
        i = 2
        while i < N:
            if not N % i:
                return False
            if i == 2:
                i -= 1
            i += 2
        return True

이렇게 하면 N이 소수인지 확인하기 위해 대략 N / 2번의 연산이 필요하게 됩니다. 그럼에도 전체 연산 횟수를 추정해 보면 (7368787 / 2) * (7368787 / 2) / 2 = 6787377731421(약 67873초)로, 시간 내에 풀 수 없습니다.

 

3번째 방법

곱셈의 핵심적인 성질이 있습니다. 어떤 수 N == a * b라고 하면, min{a, b} <= sqrt(N)입니다. 이 이유에 대해선 유리함수와 함수 y = x 사이의 관계를 생각해 보면 됩니다. N을 a로 나눠서 divisible한지 판단하나, b로 나눠서 divisible한지 판단하나 결과는 같습니다. 따라서 다음 코드가 가능합니다.

import math

def is_prime(N):
    if N == 2:
        return True
    else:
        i = 2
        while i <= math.sqrt(N):
            if not N % i:
                return False
            if i == 2:
                i -= 1
            i += 2
        return True

이러면 N을 검사하는 데 sqrt(N) / 2의 시간이 필요하게 되겠네요. 전체 연산 횟수를 추정해 보면 (sqrt(7368787) / 2) * (sqrt(7368787) / 2) / 2 = 921098(약 0.009초)로 시간 내에 여유롭게 풀 수 있습니다.

 

에라토스테네스의 체

위 방법은 사실 비효율적입니다. 이미 한 계산을 여러 번 반복하기 때문입니다. 교과서에 자주 등장하는 내용으로 에라토스테네스의 체가 있습니다. 어떻게 보면 프로그래밍의 메모이제이션(memoization) 기법을 위 방식과 합친 것이라고 볼 수 있습니다. 개념에 대해서는 검색해보면 쉽게 알 수 있으며, 직접 구현해보는 것이 도움이 많이 됩니다.

import math

prime = [True for x in range(maximum + 1)]
prime[0] = False
prime[1] = False
prime[2] = True

i = 2
while i <= math.sqrt(maximum):
    if prime[i]:
        tmp = 2 * i
        while tmp <= maximum:
            prime[tmp] = False
            tmp += i

    if i == 2:
        i -= 1
    i += 2

코드의 실행이 끝나면, prime은 소수 테이블이 됩니다. N이 소수이면 prime[N] == 1이며 N이 소수가 아니면 prime[N] == 0입니다.

 

2. 범위 지정하기

그렇다면 K번째 소수는 대략 몇 범위 안에 있을까요? 우선 2 이후로 모든 소수는 홀수이므로, 2 * K까지는 무조건 탐색해봐야 할 겁니다. 그런데 홀수이면서 소수가 아닌 수들도 고려하면 10 * K 정도로 넉넉하게 잡아야 하겠습니다. 그러나 탐색 범위를 10 * K로 하고 프로그램을 돌려 보면 10000번째 소수는 104729로 10 * 10000을 넘기 때문에 오류가 발생합니다. 여기서부턴 적절히 조정을 통해 탐색 범위를 맞추면 됩니다. 저는 30 * K를 탐색 범위로 해서 성공했습니다.

 

3. 전체 코드

이 문제는 그렇게 어렵지 않으니 코드를 보기 이전 무조건 혼자 풀어 보는 것이 좋습니다.

import math

K = int(input(''))
maximum = K * 30

prime = [True for x in range(maximum + 1)]
prime[0] = False
prime[1] = False
prime[2] = True

i = 2
while i <= math.sqrt(maximum):
    if prime[i]:
        tmp = 2 * i
        while tmp <= maximum:
            prime[tmp] = False
            tmp += i

    if i == 2:
        i -= 1
    i += 2

if K == 1:
    print(2)
else:
    j = 3
    cnt = 2

    stopCondition = False
    while not stopCondition:
        if cnt == K:
            stopCondition = True
        else:
            j += 2
            if prime[j]:
                cnt += 1

    print(j)