이전 문제 1929. 소수 구하기에서 rkddus96님의 코드를 통해 소수 구하는 빠른 방법을 배우고 여기에 적용했다. 기록을 남기자면:

from math import ceil, sqrt

for i in map(int, open(0).read().split()):
    if i == 0:
        break
    
    m, n = i, i * 2
    ls = [False] + [True] * ((n - 1) // 2)
    for x in range(1, ceil(sqrt(n) / 2)):
        if ls[x]:
            ls[2*x*x+2*x::2*x+1] = [False] * len(ls[2*x*x+2*x::2*x+1])

    if m == 2 or 2 == n:
        print(1)
    else:
        print(sum(ls[(m+1)//2:]))

그런데 성능이 메모리 36MB, 시간 216ms로 좀 더 최적화해야겠다는 생각이 들었다. 이번 문제에서도 rkddus96님의 코드가 빠른 성능을 보여줬다:

import sys
def prime(n):
    if n <= 1:
        return []
    
    sieve = [True] * (n+1)
    
    for i in range(3, int(n ** .5) + 1, 2):
        if sieve[i]:
            sieve[i*i::2*i] = [False] * len(sieve[i*i::2*i])
    
    return [2] + [i for i in range(3, n+1, 2) if sieve[i]]

def Search(prime, n):
    l,r = 0, len(prime)-1
    while l <= r:
        m = (l+r) // 2

        if prime[m] > n:
            r = m-1
        
        else:
            l = m+1
    
    return l
    
prime_list = prime(123456 * 2)

while True:
    input = int(sys.stdin.readline().strip())
    
    if input != 0:        
        print(Search(prime_list, input * 2) - Search(prime_list, input))
    
    else:
        break

메모리 약 33MB, 시간 48ms

prime(n) 함수

def prime(n):
    if n <= 1:
        return []
    
    sieve = [True] * (n+1)
    
    for i in range(3, int(n ** .5) + 1, 2):
        if sieve[i]:
            sieve[i*i::2*i] = [False] * len(sieve[i*i::2*i])
    
    return [2] + [i for i in range(3, n+1, 2) if sieve[i]]
    

1. 여기서 rkddus96님은 sieve라는 리스트를 초기화한다. 이전에 그랬던것처럼 False값으로 시작하지 않고, True값으로만 n+1개를 채운다. 여기서 n은 입력값이 아니라 아래에서 보듯 임의의 큰 수이다.
   • n+1인 이유는 먼저 리턴하는 리스트의 크기가 $n+1$개이기 때문이고, 반환값을 range를 통해 반환하는데 각 i가 인덱스와 같게 하기 위함이다(e.g. i가 3일 때 sieve[i]가 3이 되도록).
2. 밑에서 prime_list 변수를 선언하는데, prime함수를 통해 한번만 만들어놓고 저장된 값을 사용한다. 나는 반복문을 돌릴때마다 다시 만들었다.
3. 반복문의 조건이 저번과 다르게 int(n ** .5)까지 모두 사용한다. 왜일까?
   • 이번 sieve는 모든 수를 확인하기 위한 리스트이다. 저번에는 range(3, int(n**.5/2+1))이었다면 이번에는 2칸씩 건너뛰는 range로 반복문을 실행한다. 다만 반환값에서 2의 배수에 해당하는 sieve의 값이 True이더라도 3에서 2칸씩 건너뛰는 range 객체를 사용하므로 2의 배수를 소거하는 방법을 사용하고 있다. 그렇기때문에 모든 값을 가리키는 sieve라도 i가 아닌 2*i씩 건너뛰면서 소거할 수 있는 것이다.

Search(prime, n) 함수

def Search(prime, n):
    l,r = 0, len(prime)-1
    
    while l <= r:
        m = (l+r) // 2

        if prime[m] > n:
            r = m-1
        
        else:
            l = m+1
    
    return l

1. 생성된 소수 리스트 prime_list를 받는 함수.
2. prime_list의 길이에 1을 뺀 값과 0으로 r과 l을 초기화.
   • 왜 1을 뺀 값인가? 최악의 경우에도 prime[m]에서 IndexError가 나지 않게 하기 위해서
   • 그 최악의 경우는 언제인가? 만약 len(prime)을 r로 사용한다면, l이 계속 커지는 경우, 즉 n == prime[len(prime) - 1]인 경우에 IndexError가 날 것이다.
   • 그것을 조정하기 위해 else 구문에서 l = m을 대신 한다고 하자. 이때는 while문을 탈출하지 못한다. 왜냐하면 while문의 조건이 l이 r보다 같거나 작은 동안인데, r이 len(prime)으로 시작하는데, l이 m으로 되먹여지면서 새로운 m은 (l+r)를 2로 나눈 몫이기 때문이다. 예를 들어 prime[m]이 n보다 같거나 작고 r - l이 1이면 l 탈출조건을 절대 만족하지 못한다.
3. while문에서 이진탐색 알고리즘을 사용하고 있다.
   • l이 r보다 같거나 작은 동안, 가운데 있는 적당한 수 m을 구한다.
   • prime에서 인덱스가 m인 값이 입력값 n보다 큰지 확인한다. 크면 r = m-1 작거나 같으면 l = m+1을 할당하고 조정해 나간다.

시뮬레이션

prime_list = [
	2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 
    53, 59, 61, 67, 71, 73, 79, 83, 89, 97, 101, 103, 
    107, 109, 113, 127,
]   # 라 하자. (len(prime_list) == 30)

# 입력값 5에서
l, r = 0, 29

# while문 진입
m = 14
prime[14] = 47 > n
r = 13

m = 6
prime[6] = 17 > n
r = 5

m = 2
prime[2] = 5 == n
l = 3

m = 4
prime[4] = 11 > n
r = 3   # 반복문 종료

return l   # 3

자신과 같거나 작은 소수의 개수는 3개이다(2, 3, 5).
prime[m]이 n과 같아졌다고 종료하지 않는다. 유념해야겠다. 개수를 출력하는 기법.

출력

입력값이 여러개이고 마지막 값이 0이므로 0이 될 때까지 반복문을 돌리며 아래의 코드를 수행한다.

print(Search(prime_list, input * 2) - Search(prime_list, input))

입력값의 2배와 같거나 작은 소수의 개수에서 입력값과 같거나 작은 소수의 개수를 빼 출력한다.

마치면서

"소수 구하기" 문제를 통해 소수에 대해서 많이 배웠다고 생각했는데 새로운 걸 또 배웠다. 나날이 성장한다. 뿌듯하다.