[Leetcode/Python] Count Primes

Top Interview Questions 의 Easy Collection에 있는 문제입니다.
문제는 여기서 볼 수 있습니다.

 

주어진 정수 n 미만인 자연수 중 소수인 수가 몇 개인지 구하라.

 

Example 1:

Input: n = 10
Output: 4
Explanation: There are 4 prime numbers less than 10, they are 2, 3, 5, 7.

Example 2:

Input: n = 0
Output: 0

Example 3:

Input: n = 1
Output: 0

 

Constraints:

• 0 <= n <= 5 * 106

 

유의할 점은

1. n 미만인 수, 즉 0 ~ n-1 중에서 소수인 수들을 구한다는 것

2. input으로 0이 들어올 수도 있다는 점

 

위의 사항만 주의하면 문제를 푸는 것은  에라토스테네스의 체 를 활용하여 풀면 된다. "에라토스테네스의 체"라고 구글에 검색하면 위키피디아 글이 나온다. 본 글을 참고하면 쉽게 풀 수 있다.

 

Code:

class Solution:
    def countPrimes(self, n: int) -> int:
        # n 미만인 소수를 구하므로 3 미만은 소수가 존재하지 않음
        if n < 3:
            return 0
        primes = [True] * n
        primes[0] = primes[1] = False
        # 소수인지 아닌지, 2부터 int(n의 제곱근)까지 확인
        for i in range(2, int(n ** 0.5) + 1):
            # primes[i]가 소수이면
            if primes[i]:
                # primes[i]의 제곱부터 시작해 primes[i]의 배수를 모두 False처리
                # primes[i]*2부터 시작하지 않는 이유 : 그 전에 i보다 작은 소수의 차례일 때 이미 false처리 했기 때문
                primes[i**2:n:i] = [False] * len(primes[i**2:n:i])
        return sum(primes)

 

코드 내에서 유의할 점이 한 가지 있는데, for문 안의 if문이다. 여기서 i가 소수일 경우, 범위 내(즉 n 미만의 수)의 자연수에 한하여 i를 제외한 i의 배수를 모두 False처리(즉 i의 배수는 소수가 아니라는 의미)하는 것이 정석이다.

그런데 위의 코드에서는 i의 배수를 모두 False처리하되, i의 제곱부터 시작하고 있다.

그 이유는 예를 들어 i가 5라고 할 때, for문에서 5보다 작은 소수일 경우(2,3)에서 모두 이미 False처리했기 때문에 굳이 재차 함으로써 시간복잡도를 높일 필요가 없기 때문이다.

 

위의 코드처럼 i의 제곱부터 False처리를 함으로써 코드가 더 빨리 돌아가게 할 수 있다.