[Python] 소수 판별 알고리즘

난이도가 낮은 알고리즘 문제 중에서 자주 나오는 문제가 소수 판별 문제이다.

소수 : 1과 자기 자신 이외에 약수를 가지지 않는 2 이상의 자연수

소수를 판별하는 방법은 크게 두 가지로 나뉜다.

반복문 사용

첫 번째 방법은 반복문 사용이다. 자연수 n이 소수인지 판별하기 위해서 2부터 \( [\sqrt{n}] \)까지 나누었을 때 나머지가 0이 있는지 확인하는 방법이다.

Code

from math import sqrt

def is_prime(n):
    if n < 2:
        return False
    for i in range(2, int(sqrt(n)) + 1):
        if n % i == 0:
            return False
    return True

sqrt(n)까지만 확인하는 이유는 \( \sqrt{n} \) 보다 큰 약수는 이미 작은 약수를 확인하면서 확인하기 때문이다. 예를 들어서 n이 20인 경우 36의 약수는 1, 2, 3, 4, 6, 9, 12, 18, 36 이지만 9, 12, 18은 4, 3, 2와 한쌍이기 때문에 6까지만 검사하게 되면 9, 12, 18이 36의 약수인 것을 알 수 있다.

에라토스테네스의 체

아스토스테네스의 체는 여러 개의 소수를 판별할 때 유용하다. 또는 큰 수에 대해서 소수 여부를 판별할 때 시간복잡도가 많이 개선된다. 크기가 n인 배열을 만든 후에 처음 두 칸은 False로 설정하고 이후의 배열은 모두 True로 초기화해준다. 이후 배열을 한 번 순회하면서 소수를 찾고 소수의 배수들을 모두 False로 변환해준다.

Code

def is_prime(n):
    prime = [False, False, True] + [True, False] * ( (n - 1) // 2)
    # 2를 제외한 모든 짝수는 소수가 아님
    for i in range(3, n+1, 2):
        if prime[i]:
            for j in range(i*i, n+1, 2*i):
                prime[j] = False

    return prime[n]

코드의 베이스라인은 위와 같다. 소수 판별 문제에서는 어떤 수가 소수인지 판별하는 문제도 있지만 n이하의 자연수 중 소수가 몇 개인지, 어떤 수들인지를 알기 위한 문제가 많다. j에 대한 for문에서 범위가 i * i 부터 n + 1까지 간격은 2 * i로 큰 수들을 False로 변환한다.

i * i부터 시작하는 이유는 i * (i - 2), i * (i - 4), ...은 i - 2번 째 반복문에서 (i - 2) * (i - 2) + 2 * (i - 2)가 되어서 한 번 방문을 했거나 또는 소수가 아닌 수의 배수이기 때문에 방문할 필요가 없다.

2 * i씩 건너띄는 이유는 간단하다. 홀수에 짝수를 더하면 홀수이기 때문이다. 짝수는 2를 제외하고는 소수가 아니기 때문에 굳이 방문할 필요가 없다.

n = 1000000
prime = [False, False, True] + [True, False] * ((n - 1) // 2)
prime_number = [2]
for i in range(3, n+1, 2):
    if prime[i]:
        if i > 1000:                    
            prime_number.append(i)
        for j in range(i*i, n+1, 2*i):    # i * i 이하의 수는 i 전에 한 번씩 검사됨
            prime[j] = False

print(prime_number)        # 1,000과 1,000,000 사이의 소수
print(len(prime_number))# 1,000과 1,000,000 사이 소수의 개수

위에서 처럼 "1,000보다 큰 소수를 모두 구하여라." 같은 문제를 풀 때 에라토스테네스의 체가 유용하게 쓰이고 많은 알고리즘 문제를 해결할 때 에라토스테네스의 체를 사용하다.
하지만 무조건적으로 에라토스 테네스의 체가 좋다고 얘기할 순 없다. n이 소수인지 판별하기 위해서 크기가 n+1인 리스트를 생성해야 한다. 이론상 O(Nlog(logN))이라고 되어 있고 반복문을 사용할 때는 O(sqrt(N))이다.

1000000007 과 같은 큰 소수가 소수인지 판별하는 데에는 반복문을 사용하는 것이 훨씬 빠르다는 얘기다.

Code

import time
from math import sqrt


def is_prime(n):
    prime = [False, False] + [True] * (n - 1)
    for i in range(3, int(sqrt(n))+1, 2):
        if prime[i]:
            for j in range(i*i, n+1, 2*i):
                prime[j] = False

    return prime[n]


def is_prime2(n):
    if n % 2 == 0:
        return False
    for i in range(3, int(sqrt(n))+1, 2):
        if n % i == 0:
            return False
    return True


start = time.time()
is_prime(1000000007)        # 에라토스테네스의 체
print(time.time() - start)

start = time.time()
is_prime2(1000000007)       # 반복문 사용
print(time.time() - start)

에라토스테네스의 체도 sqrt(n)까지만 해서 n이 소수인지 판별하도록 했음에도 불구하고 출력 결과에서 보면 많이 차이가 나는 것을 알 수 있다.

결론

에라토스테네스의 체를 사용하게 되면 여러 개의 소수를 한 번에 구할 수 있고, 한 번 구해놓고 나면 O(1)로도 판별이 가능하기 때문에 여러 개의 소수를 구하는 경우에는 꼭 에라토스테네스의 체를 사용하는 것이 좋다. 입력이 하나이고 그 수가 소수인지 판별할 때는 그냥 간단하게 반복문을 사용하는 것이 좋다.