Number-theoretic transform in Python
FFT over a finite field Z_p with a primitive root of unity, giving exact convolution modulo p.
- File
number_theoretic_transform.py, 100 lines- Level
- I0 not checked yet
- Toolchain
- Python 3.14, pytest 9
- Test framework
- pytest
- Run the tests
python3 -m pytest
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/number-theoretic-transform/. The authors are listed in prov/7AWXWS.jsonl.
# Primes p = c * 2^k + 1 with the primitive root 3, as (p, g, k).
NTT_PRIMES = [(998244353, 3, 23), (167772161, 3, 25), (469762049, 3, 26)]
def root_of_unity(n: int, p: int, g: int) -> int:
"""Return g^((p-1)/n) mod p, a primitive n-th root of unity modulo p.
Precondition: p is prime, g is a primitive root modulo p, and n divides
p - 1; otherwise ValueError.
"""
if n < 1 or (p - 1) % n:
raise ValueError("n must divide p - 1")
return pow(g, (p - 1) // n, p) # @line L4
def ntt(x: list[int], w: int, p: int) -> None:
"""Replace x by X[k] = sum of x[j] * w^(j*k) mod p, in place.
Precondition: n = len(x) is a power of two, p is an odd prime, and w has
order n modulo p. Cost: n lg n additions and (n/2) lg n multiplications
modulo p, n/2 - 1 more for the table, and O(n) index steps.
"""
n = len(x)
if n < 1 or n & (n - 1):
raise ValueError("the length must be a power of two")
t = n.bit_length() - 1
rev = [0] * n
for j in range(1, n):
rev[j] = (rev[j >> 1] >> 1) | ((j & 1) << (t - 1))
for j in range(n): # @line L10
if j < rev[j]:
x[j], x[rev[j]] = x[rev[j]], x[j]
table = [1] * max(n // 2, 1) # @line L11
for k in range(1, n // 2):
table[k] = table[k - 1] * w % p
h = 1 # @line L12
while h < n: # @line L13
step = n // (2 * h)
for s in range(0, n, 2 * h): # @line L14
for k in range(h): # @line L15
u = x[s + k] # @line L16
v = table[k * step] * x[s + k + h] % p # @line L16
x[s + k] = (u + v) % p # @line L17
x[s + k + h] = (u - v) % p # @line L17
h *= 2 # @line L18
def multiply_mod(a: list[int], b: list[int], p: int = 998244353, g: int = 3) -> list[int]:
"""Return the coefficients of a(t) * b(t) modulo the prime p.
Precondition: g is a primitive root modulo p, and the least power of two
n >= len(a) + len(b) - 1 divides p - 1. Cost: three NTTs of length n,
n pointwise products and n scalings, so O(n log n) operations modulo p.
"""
if not a or not b:
return []
m = len(a) + len(b) - 1 # @line L1
n = 1 # @line L2
while n < m:
n *= 2
if (p - 1) % n: # @line L3
raise ValueError("p - 1 is not divisible by the transform length")
w = root_of_unity(n, p, g)
fa = [v % p for v in a] + [0] * (n - len(a)) # @line L5
fb = [v % p for v in b] + [0] * (n - len(b)) # @line L5
ntt(fa, w, p) # @line L6
ntt(fb, w, p) # @line L6
fc = [u * v % p for u, v in zip(fa, fb)] # @line L7
ntt(fc, pow(w, -1, p), p) # @line L8
n_inv = pow(n, -1, p)
return [fc[j] * n_inv % p for j in range(m)] # @line L9
def multiply_exact(a: list[int], b: list[int]) -> list[int]:
"""Return the exact integer coefficients of a(t) * b(t).
Uses the three primes of NTT_PRIMES and the Chinese remainder theorem.
Precondition: every coefficient c of the product has |c| < P/2, where P
is the product of the three primes (about 7.9 * 10^25); this holds when
min(len(a), len(b)) * max|a| * max|b| < P/2. Length limit: 2^23.
"""
if not a or not b:
return []
res = [multiply_mod(a, b, p, g) for p, g, _ in NTT_PRIMES]
(p1, _, _), (p2, _, _), (p3, _, _) = NTT_PRIMES
big = p1 * p2 * p3
inv1 = pow(p1, -1, p2)
inv12 = pow(p1 * p2, -1, p3)
out = []
for r1, r2, r3 in zip(*res):
# Garner's mixed-radix form: c = r1 + p1 * (y2 + p2 * y3).
y2 = (r2 - r1) * inv1 % p2
c12 = r1 + p1 * y2
y3 = (r3 - c12) * inv12 % p3
c = c12 + p1 * p2 * y3
out.append(c - big if c > big // 2 else c)
return out
test_number_theoretic_transform.py 130 lines
# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/number-theoretic-transform/. The authors are listed in prov/7AWXWS.jsonl.
import random
import pytest
from number_theoretic_transform import NTT_PRIMES, multiply_exact, multiply_mod, ntt, root_of_unity
# (prime, primitive root): 3 for 17 and 998244353, 5 for 97.
PRIMES = [(17, 3), (97, 5), (998244353, 3)]
def naive_dft(x: list[int], w: int, p: int) -> list[int]:
n = len(x)
return [sum(x[j] * pow(w, j * k, p) for j in range(n)) % p for k in range(n)]
def schoolbook(a: list[int], b: list[int], p: int | None = None) -> list[int]:
c = [0] * (len(a) + len(b) - 1)
for i, u in enumerate(a):
for j, v in enumerate(b):
c[i + j] += u * v
return c if p is None else [v % p for v in c]
def max_len(p: int) -> int:
n = 1
while (p - 1) % (2 * n) == 0 and n < 64:
n *= 2
return n
def test_roots() -> None:
for p, g in PRIMES + [(q, h) for q, h, _ in NTT_PRIMES]:
n = max_len(p)
while n >= 1:
w = root_of_unity(n, p, g)
assert pow(w, n, p) == 1
if n > 1:
assert pow(w, n // 2, p) == p - 1
n //= 2
with pytest.raises(ValueError):
root_of_unity(3, 17, 3)
def test_examples() -> None:
assert root_of_unity(4, 17, 3) == 13
x = [1, 2, 3, 0]
ntt(x, 13, 17)
assert x == [6, 7, 2, 6]
assert multiply_mod([1, 2, 3], [4, 5], 17) == [4, 13, 5, 15]
assert multiply_mod([1, 2, 3], [4, 5]) == [4, 13, 22, 15]
assert multiply_exact([1, 2, 3], [4, 5]) == [4, 13, 22, 15]
def test_edges() -> None:
for p, g in PRIMES:
for n in (1, 2, 4, 8, 16):
w = root_of_unity(n, p, g)
zeros = [0] * n
ntt(zeros, w, p)
assert zeros == [0] * n
for pos in range(n):
imp = [int(j == pos) for j in range(n)]
ntt(imp, w, p)
assert imp == [pow(w, pos * k, p) for k in range(n)]
const = [3] * n
ntt(const, w, p)
assert const == [3 * n % p] + [0] * (n - 1)
assert multiply_mod([5], [7], p, g) == [35 % p]
assert multiply_mod([0, 0, 0], [1, 2], p, g) == [0, 0, 0, 0]
assert multiply_mod([1], [3, 1, 4], p, g) == [3, 1, 4]
assert multiply_mod([], [1, 2]) == []
assert multiply_exact([2, -3], [-4, 5]) == [-8, 22, -15]
def test_errors() -> None:
for n in (0, 3, 6, 12):
with pytest.raises(ValueError):
ntt([1] * n, 1, 17)
# Length 32 does not divide 16 = 17 - 1.
with pytest.raises(ValueError):
multiply_mod([1] * 10, [1] * 10, 17, 3)
def test_random_against_naive() -> None:
rng = random.Random(1971)
for _ in range(500):
p, g = PRIMES[rng.randrange(len(PRIMES))]
n = 1 << rng.randint(0, max_len(p).bit_length() - 1)
w = root_of_unity(n, p, g)
x = [rng.randrange(p) for _ in range(n)]
y = list(x)
ntt(y, w, p)
assert y == naive_dft(x, w, p)
ntt(y, pow(w, -1, p), p)
n_inv = pow(n, -1, p)
assert [v * n_inv % p for v in y] == x
def test_random_multiply() -> None:
rng = random.Random(1975)
for _ in range(200):
p, g = PRIMES[rng.randrange(len(PRIMES))]
lim = max_len(p)
la = rng.randint(1, lim // 2)
lb = rng.randint(1, lim // 2)
a = [rng.randrange(p) for _ in range(la)]
b = [rng.randrange(p) for _ in range(lb)]
assert multiply_mod(a, b, p, g) == schoolbook(a, b, p)
def test_exact_large_coefficients() -> None:
rng = random.Random(1959)
for _ in range(30):
la, lb = rng.randint(1, 120), rng.randint(1, 120)
a = [rng.randint(-(10**9), 10**9) for _ in range(la)]
b = [rng.randint(-(10**9), 10**9) for _ in range(lb)]
assert multiply_exact(a, b) == schoolbook(a, b)
def test_large_length() -> None:
rng = random.Random(23)
a = [rng.randrange(800) for _ in range(1500)]
b = [rng.randrange(800) for _ in range(1500)] # 1500 * 800^2 < p, so exact
c = multiply_mod(a, b)
for m in rng.sample(range(len(c)), 10):
ref = sum(a[i] * b[m - i] for i in range(max(0, m - 1499), min(m, 1499) + 1))
assert c[m] == ref
What the next level needs
- the tests did not run in this build