[백준/Python] 4948: 베르트랑 공준

[백준/Python] 4948: 베르트랑 공준

임의의 자연수 n이 주어질 때 n보다 크고 2n보다 작거나 같은 소수의 개수를 찾는 문제이다.

처음엔 간단한 문제라고 생각해서 1929: 소수 구하기의 코드를 약간 변형해서 그대로 제출했는데 시간 초과가 났다. (시간 제한은 1초)

뭐가 문제일까 고민해보다가 입력이 반복해서 주어질 때 많은 중복 계산이 일어난다는 것을 깨달았다. 그래서 Memoization을 사용하는 동적 계획법 방식을 채택해 보았다.
메모이제이션이란 동일한 계산을 반복할 때 이미 계산된 값을 저장해 두고 다음 계산때 또 사용하는 것이다.

테스트 케이스에선 여러 개의 입력이 주어진다. 입력마다 그 숫자가 소수인지 아닌지를 계속해서 판별하는것은 큰 시간 낭비이다. 모든 소수 판별의 경우의 수는 $123,456*2$이다. 배열을 우선 False로 초기화한 후 isPrime 함수를 통해 처음에 소수 판별을 먼저 다 해놓는다.
그 후에 다시 for문을 돌면서 n의 입력이 들어오면 카운트를 세고 출력하였다.

소스코드

---

def isPrime(n):
    if n == 1:
        return False
    for i in range(2, int(n**(1/2)) + 1):
        if n%i == 0:
            return False
    return True

arr = [False]*(123456*2+1)

for i in range(len(arr)):
    if isPrime(i) is True:
        arr[i] = True

#

n = int(input())
while n != 0:
    cnt = 0
    for i in range(n+1, 2*n+1):
        if arr[i] is True:
            cnt += 1
    print(cnt)
    n = int(input())

그런데 통과는 했지만 개선이 필요해 보인다. 시간이 1056ms나 걸리고 있기 때문이다. 다른 언어도 아니고 파이썬으로 푸는데 효율적인 알고리즘을 못만들어서 겁나 오래 걸리는건 뭔가 자존심 상하기도 함..

방법을 찾아보다가 에라토스테네스의 체라는 알고리즘을 찾았다. 역시 위의 메모이제이션 방식과 비슷하다.

개선 소스코드

---

arr = [True]*(123456*2+1)
arr[1] = False

m = int(123456*2**(1/2))
for i in range(2, m+1):
    if arr[i] == True: 
        for j in range(i+i, 123456*2+1, i): 
            arr[j] = False

n = int(input())
while n != 0:
    cnt = sum(arr[n+1:2*n+1])
    print(cnt)
    n = int(input())

소수는 1과 자기 자신만을 약수로 갖는 수이다. 배열을 전부 True로 초기화한 후 2부터 시작했을 때 True인 어떤 수를 처음 만나면 그 수는 당연히 소수이므로 그 수의 배수를 모두 False로 만든다. 또 루프를 돌다가 True인 수를 만나면 그 수는 무조건 소수이고, 그 수의 배수를 모두 False로 만든다.

배열에 끝에 도달할 때까지 이를 반복하면 배열의 인덱스가 소수인지 아닌지를 나타내는 배열이 만들어진다.

n을 입력받은 후에 n보다 크고 2n보다 작거나 같은 소수 개수를 찾을 때, 반복문을 돌지 않고 배열의 True개수를 찾아 합해주는 sum()을 사용하면 더 간편하게 문제를 풀 수 있었다.