Divisor-count sieve in Python

Compute the number of divisors d(m) for all m up to n by adding 1 at every multiple of each e from 1 to n. It takes Θ(n log n) word operations, and adding e instead of 1 gives sigma(m).

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

divisor_count_sieve.pyRaw GitHub

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


def divisor_count_sieve(n: int) -> list[int]:
    """Return d[0 .. n], where d[m] is the number of positive divisors of m for m >= 1, and d[0] = 0.

    Precondition: n >= 0.
    Cost: Theta(n log n) word operations; the output list is the only table.
    """
    d = [0] * (n + 1)  # @line L1
    for e in range(1, n + 1):  # @line L2
        for m in range(e, n + 1, e):  # @line L3
            d[m] += 1  # @line L4
    return d  # @line L5


def divisor_power_sum_sieve(n: int, k: int = 1) -> list[int]:
    """Return s[0 .. n], where s[m] = sigma_k(m), the sum of e**k over the divisors e of m, and s[0] = 0.

    Precondition: n >= 0 and k >= 0.
    Cost: Theta(n log n) additions on numbers up to about n**k times log log n.
    """
    s = [0] * (n + 1)
    for e in range(1, n + 1):
        w = e**k
        for m in range(e, n + 1, e):
            s[m] += w
    return s

test_divisor_count_sieve.py 70 lines

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

import random

from divisor_count_sieve import divisor_count_sieve, divisor_power_sum_sieve


def divisors_by_definition(m: int) -> list[int]:
    return [e for e in range(1, m + 1) if m % e == 0]


def divisors_by_trial_division(m: int) -> list[int]:
    small, large, e = [], [], 1
    while e * e <= m:
        if m % e == 0:
            small.append(e)
            if e * e != m:
                large.append(m // e)
        e += 1
    return small + large[::-1]


def test_example() -> None:
    assert divisor_count_sieve(12) == [0, 1, 2, 2, 3, 2, 4, 2, 4, 3, 4, 2, 6]
    assert divisor_power_sum_sieve(12) == [0, 1, 3, 4, 7, 6, 12, 8, 15, 13, 18, 12, 28]
    assert sum(divisor_count_sieve(12)) == 35


def test_edges() -> None:
    assert divisor_count_sieve(0) == [0]
    assert divisor_count_sieve(1) == [0, 1]
    assert divisor_count_sieve(2) == [0, 1, 2]
    assert divisor_power_sum_sieve(0, 2) == [0]
    assert divisor_power_sum_sieve(1, 0) == [0, 1]


def test_against_definition_up_to_800() -> None:
    d = divisor_count_sieve(800)
    s1 = divisor_power_sum_sieve(800, 1)
    s2 = divisor_power_sum_sieve(800, 2)
    for m in range(1, 801):
        divs = divisors_by_definition(m)
        assert d[m] == len(divs)
        assert s1[m] == sum(divs)
        assert s2[m] == sum(e * e for e in divs)


def test_every_n_up_to_2000() -> None:
    full = [0] + [len(divisors_by_trial_division(m)) for m in range(1, 2001)]
    for n in range(0, 2001, 7):
        assert divisor_count_sieve(n) == full[: n + 1]
    for n in range(1990, 2001):
        assert divisor_count_sieve(n) == full[: n + 1]


def test_random_larger_n() -> None:
    rng = random.Random(1849)
    for _ in range(4):
        n = rng.randint(2001, 50000)
        d = divisor_count_sieve(n)
        s = divisor_power_sum_sieve(n)
        for m in rng.sample(range(1, n + 1), 300) + [n]:
            divs = divisors_by_trial_division(m)
            assert d[m] == len(divs)
            assert s[m] == sum(divs)


def test_summatory_value() -> None:
    assert sum(divisor_count_sieve(10**4)) == 93668

What the next level needs

the tests did not run in this build

Read the full page about Divisor-count sieve