Fast Python Prime Counting using Optimised Sieve of Eratosthenes

Harnessing the slice ability of Python

Sieve of Eratosthenes is a simple and ancient algorithm used to find the prime numbers. It is one of the most efficient ways to find small prime numbers.

For a given upper limit n, the algorithm works by iteratively marking the multiples of primes as composite, starting from 2. Once all multiples of 2 have been marked composite, the muliples of next prime, ie 3 are marked composite. This process continues until p ≤= sqrt(n) where p is a prime number.

Algorithm to implement in Python is as such:

def countPrimes(self, n: int) -> int:
    if n <= 2:
        return 0
    sieve = [1]*(n) # 0, 1, 2, ... n-1, exclude n since find primes          
                    # less than n; true means prime
    sieve[0] = 0
    sieve[1] = 0
    p = 2 # first prime
    while p*p <= n:
        for i in range(p*p, n):    
            if i % p == 0: # multiples of p, 
                       # so make as composites (not primes)
                sieve[i] = 0
        for i in range(p+1, n): # find the next prime p
            if sieve[i]:
                p = i
                break
    return sum(sieve)

While the above is a faithful reproduction of the algorithm, it is slow. If you look at the “While” loop, it’s basically execution with O(n*n) time. Hence, if you attempt the question on leetcode of counting prime, you will have time limit error. Their testcases have very high value of n. So even if you know how to code the Sieve of Eratosthenes, you still need to optimise the run-time.

Python does have the slicing features. We oft extract a shallow copy of an iterable list using array[start:end:step] syntax. But we seldom using it to change out the values of the slice.

Focusing on the section of the code to optimise:

for p in range(2, int(n**0.5) + 1):
    if sieve[p]: # is a prime
        # multiples of p is marked as 0 i.e. composite
        sieve[p*p:n:p] = [0]*len(sieve[p*p:n:p])      

The above optimised code of marking composite starting from the first mutiple of p (after p itself) which is p*p to end of sieve array, and finding the next prime is fast at O(n) i.e. linear time.