[LeetCode/Python] 204. Count Primes

https://leetcode.com/problems/count-primes/

 

Step1. n개의 True로 구성된 is_prime 리스트 생성

Step2. 0과 1에 대해 소수가 아니므로 False로 미리 처리

Step3. 2부터 루트 n까지 반복하면서

Step4. 그 수의 배수들을 False 처리

Step5. is_prime 리스트에서 최종적으로 True의 개수만 세기

 

import math

class Solution(object):
    def countPrimes(self, n):
        """
        :type n: int
        :rtype: int
        """
        if n == 0 or n == 1:
            return 0

        is_prime = [True] * (n)
        is_prime[0] = is_prime[1] = False
        
        root_n = int(math.sqrt(n))

        for i in range(2, root_n+1):
            # multiple numbers
            if is_prime[i]:
                for j in range(i*i, n, i):
                    is_prime[j] = False
                
        count = 0

        for flag in is_prime:
            if flag == True:
                count += 1
        
        return count