제곱근을 이용한 소수 판별

소수는 1보다 큰 자연수 중 약수¹가 1과 자기 자신뿐인 수다.
코딩 테스트에서 어떤 수 $n$이 소수인지 확인해야 할 때, 2부터 $n - 1$까지 모두 나누어보는 방법을 먼저 떠올릴 수 있다.

하지만 약수는 항상 짝을 이룬다. 이 성질을 이용하면 $\sqrt{n}$까지만 확인해도 소수 여부를 판별할 수 있다.

왜 $\sqrt{n}$까지만 확인할까?

약수는 서로 곱해 원래 수가 되는 두 수의 쌍으로 나타낼 수 있다. 36의 약수 쌍을 살펴보자.

1 × 36
2 × 18
3 × 12
4 × 9
6 × 6

$\sqrt{36} = 6$이고, 각 약수 쌍에는 6 이하인 수가 적어도 하나 있다.
두 수가 모두 6보다 크다면 두 수의 곱도 36보다 커지므로 36의 약수 쌍이 될 수 없다.

일반적인 수 $n$도 마찬가지다. $\sqrt{n}$보다 큰 약수가 있다면 그와 짝을 이루는 약수는 반드시 $\sqrt{n}$보다 작다. 따라서 2부터 $\sqrt{n}$까지 나누어떨어지는 수가 없다면, 그보다 큰 범위에도 약수가 없다.

이 원리를 코드로 옮기면 다음과 같다.

def is_prime(n):
    if n < 2:
        return False

    i = 2
    while i * i <= n:
        if n % i == 0:
            return False
        i += 1

    return True

assert is_prime(2)
assert not is_prime(1)
assert not is_prime(9)

0과 1은 소수가 아니므로 n < 2일 때는 바로 False를 반환한다.
i * i <= n 조건으로 $\sqrt{n}$ 이하의 약수만 확인한다.

전체 수를 확인하는 대신 최대 $\sqrt{n}$번만 나누어보므로 시간 복잡도는 $O(\sqrt{n})$이다.