LeetCode 204. Count Primes

September 13, 2020

We can solve this naively in $O(n^2)$ time by:

Step through each integer smaller than $n$ and factorize each.
Increment a counter for each that is prime.

That's a pretty dumb solution, so I googled it and found the Sieve of Erathosthenes algorithm.

We'll use that instead.

import math
import unittest

class Test(unittest.TestCase):

    def setUp(self):
        self.s = Solution()

    def test_examples(self):
        self.assertEqual(self.s.countPrimes(10), 4)
        self.assertEqual(self.s.countPrimes(2), 0)

    def test_zero(self):
        self.assertEqual(self.s.countPrimes(0), 0)

    def test_prime_boundary(self):
        self.assertEqual(self.s.countPrimes(3), 1)
        self.assertEqual(self.s.countPrimes(5), 2)

    def test_larger(self):
        self.assertEqual(self.s.countPrimes(121), 30)

class Solution:

    def countPrimes(self, n: int) -> int:
        if n <= 2:
            return 0
        # the number x is represented by a[x]
        a = [1] * n
        # ignore 0 and 1
        a[0] = a[1] = 0
        nsqrt = math.ceil(math.sqrt(n))
        for i in range(2, nsqrt + 1):
            if not a[i]:
                continue
            j = int(math.pow(i, 2))
            while j < n:
                a[j] = 0
                j += i
        return sum(a)

if __name__ == "__main__":
    unittest.main()