Sieve of Eratosthenes in Python

List all primes up to n. Keep a table of flags, and for each prime p with p^2 at most n clear the flags of p^2, p^2 + p, p^2 + 2p, and so on. It takes Θ(n log log n) word operations and n + 1 flags.

File: sieve_of_eratosthenes.py, 33 lines
Level: I0 not checked yet
Toolchain: Python 3.14, pytest 9
Test framework: pytest
Run the tests: python3 -m pytest

sieve_of_eratosthenes.pyRaw GitHub

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/sieve-of-eratosthenes/. The authors are listed in prov/F0W4PP.jsonl.


def sieve_of_eratosthenes(n: int) -> list[int]:
    """Return the increasing list of all primes p <= n.

    Precondition: n is an integer (n < 2 gives the empty list).
    Cost: Theta(n log log n) word operations and n + 1 bytes of flags.
    """
    if n < 2:  # @line L1
        return []
    is_prime = bytearray([1]) * (n + 1)  # @line L2
    is_prime[0] = 0
    is_prime[1] = 0
    p = 2  # @line L3
    while p * p <= n:  # @line L4
        if is_prime[p]:  # @line L5
            for m in range(p * p, n + 1, p):  # @line L6
                is_prime[m] = 0  # @line L7
        p += 1  # @line L8
    return [m for m in range(2, n + 1) if is_prime[m]]  # @line L9


def prime_flags(n: int) -> bytearray:
    """Return a table T of n + 1 bytes (empty for n < 0) with T[m] = 1 exactly when m is prime.

    Precondition: n is an integer. Cost: Theta(n log log n) word operations.
    """
    table = bytearray(max(n + 1, 0))
    for p in sieve_of_eratosthenes(n):
        table[p] = 1
    return table

test_sieve_of_eratosthenes.py 67 lines

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/sieve-of-eratosthenes/. The authors are listed in prov/F0W4PP.jsonl.

import random

from sieve_of_eratosthenes import prime_flags, sieve_of_eratosthenes


def is_prime_by_trial_division(m: int) -> bool:
    if m < 2:
        return False
    d = 2
    while d * d <= m:
        if m % d == 0:
            return False
        d += 1
    return True


LIMIT = 40000
ORACLE = [m for m in range(LIMIT + 1) if is_prime_by_trial_division(m)]


def oracle_primes(n: int) -> list[int]:
    return [p for p in ORACLE if p <= n]


def test_example() -> None:
    assert sieve_of_eratosthenes(30) == [2, 3, 5, 7, 11, 13, 17, 19, 23, 29]
    assert len(sieve_of_eratosthenes(100)) == 25


def test_edges() -> None:
    assert sieve_of_eratosthenes(-5) == []
    assert sieve_of_eratosthenes(0) == []
    assert sieve_of_eratosthenes(1) == []
    assert sieve_of_eratosthenes(2) == [2]
    assert sieve_of_eratosthenes(3) == [2, 3]
    assert sieve_of_eratosthenes(4) == [2, 3]
    # n is a prime, a square of a prime, or one less than one
    for n in (7, 8, 9, 24, 25, 26, 48, 49, 50, 120, 121, 122, 168, 169, 170):
        assert sieve_of_eratosthenes(n) == oracle_primes(n)


def test_every_n_up_to_2000() -> None:
    for n in range(2001):
        assert sieve_of_eratosthenes(n) == oracle_primes(n)


def test_random_larger_n() -> None:
    rng = random.Random(4127)
    for _ in range(20):
        n = rng.randint(2001, LIMIT)
        assert sieve_of_eratosthenes(n) == oracle_primes(n)


def test_known_counts() -> None:
    assert len(sieve_of_eratosthenes(10**4)) == 1229
    assert len(sieve_of_eratosthenes(10**6)) == 78498


def test_flags() -> None:
    for n in (-1, 0, 1, 2, 30, 1000):
        table = prime_flags(n)
        assert len(table) == max(n + 1, 0)
        for m in range(len(table)):
            assert table[m] == int(is_prime_by_trial_division(m))

What the next level needs

the tests did not run in this build

Read the full page about Sieve of Eratosthenes