FFT-based polynomial multiplication in Python

# SPDX-License-Identifier: MIT
# Part of https://algorithm.tamnd.com/fft-polynomial-multiplication/. The authors are listed in prov/WT2CHV.jsonl.

"""Polynomial multiplication by the fast Fourier transform.

Polynomials are coefficient lists, lowest degree first, normalized (no zero
on top; the zero polynomial is []). The exact transform works over Z/pZ for
a prime p with 2^k | p - 1 (a number-theoretic transform). Products over Z,
over Z/mZ for any m, and over Q are reduced to it by several primes and the
Chinese remainder theorem. A floating point version over C is included for
comparison. Costs are in ring operations of Z/pZ: O(N log N) for N >= m + n - 1.
"""

from __future__ import annotations

import cmath
from fractions import Fraction
from math import lcm
from typing import List, Sequence

# Primes p = c 2^k + 1 with a primitive root g, for the exact transform.
NTT_PRIMES = [(998244353, 3), (167772161, 3), (469762049, 3), (754974721, 11)]


def _normalize(c: List) -> List:
    while c and c[-1] == 0:
        c.pop()
    return c


def ntt(a: List[int], omega: int, p: int) -> None:
    """In-place DFT of a, of length N = 2^k, at the powers of omega modulo p.

    Precondition: omega^N = 1 and omega^(N/2) = -1 modulo p.
    Cost: (N/2) log2 N butterflies, each 2 multiplications and 2 additions.
    """
    n = len(a)
    j = 0
    for i in range(1, n):  # @line F1: bit-reversal permutation
        bit = n >> 1
        while j & bit:
            j ^= bit
            bit >>= 1
        j |= bit
        if i < j:
            a[i], a[j] = a[j], a[i]
    length = 2  # @line F2
    while length <= n:  # @line F3
        w_len = pow(omega, n // length, p)  # @line F4
        half = length // 2
        for s in range(0, n, length):  # @line F5
            w = 1  # @line F6
            for k in range(half):  # @line F7
                t = w * a[s + k + half] % p  # @line F8
                x = a[s + k]
                a[s + k] = (x + t) % p  # @line F9
                a[s + k + half] = (x - t) % p
                w = w * w_len % p  # @line F10
        length <<= 1  # @line F11


def multiply_mod(a: Sequence[int], b: Sequence[int], p: int = 998244353, g: int = 3) -> List[int]:
    """The product a b over Z/pZ by the number-theoretic transform.

    Precondition: p is prime, g generates (Z/pZ)^*, and 2^k divides p - 1 for
    the power of two N >= len(a) + len(b) - 1. Cost: O(N log N) operations mod p.
    """
    m, n = len(a), len(b)
    if m == 0 or n == 0:  # @line L1
        return []
    size = 1
    while size < m + n - 1:  # @line L2
        size <<= 1
    if (p - 1) % size:
        raise ValueError("p - 1 is not divisible by the transform length")
    fa = [x % p for x in a] + [0] * (size - m)  # @line L3
    fb = [x % p for x in b] + [0] * (size - n)
    omega = pow(g, (p - 1) // size, p)
    ntt(fa, omega, p)  # @line L4
    ntt(fb, omega, p)
    for i in range(size):  # @line L5
        fa[i] = fa[i] * fb[i] % p
    ntt(fa, pow(omega, p - 2, p), p)  # @line L6
    inv = pow(size, p - 2, p)
    c = [x * inv % p for x in fa[: m + n - 1]]
    return _normalize(c)  # @line L7


def _is_prime(n: int) -> bool:
    """Deterministic Miller-Rabin for n < 3.4 * 10^14."""
    if n < 2:
        return False
    for q in (2, 3, 5, 7, 11, 13, 17):
        if n % q == 0:
            return n == q
    d, s = n - 1, 0
    while d % 2 == 0:
        d, s = d // 2, s + 1
    for b in (2, 3, 5, 7, 11, 13, 17):
        x = pow(b, d, n)
        if x in (1, n - 1):
            continue
        for _ in range(s - 1):
            x = x * x % n
            if x == n - 1:
                break
        else:
            return False
    return True


def _primitive_root(p: int) -> int:
    """The least generator of (Z/pZ)^* for a prime p with p - 1 = c 2^k, c small."""
    m, factors, q = p - 1, [2], 3
    while m % 2 == 0:
        m //= 2
    while q * q <= m:
        if m % q == 0:
            factors.append(q)
            while m % q == 0:
                m //= q
        q += 2
    if m > 1:
        factors.append(m)
    g = 2
    while any(pow(g, (p - 1) // f, p) == 1 for f in factors):
        g += 1
    return g


MAX_LOG = 20  # the extra primes support transform lengths up to 2^20


def _crt_primes(count: int, size: int) -> List[tuple]:
    """count primes p with size | p - 1, each with a primitive root."""
    if size > 1 << MAX_LOG:
        raise ValueError("transform length too large")
    out = [(p, g) for p, g in NTT_PRIMES if (p - 1) % size == 0][:count]
    c = 1
    while len(out) < count:
        p = (c << MAX_LOG) + 1
        if _is_prime(p) and all(p != q for q, _ in out):
            out.append((p, _primitive_root(p)))
        c += 1
    return out


def multiply_int(a: Sequence[int], b: Sequence[int]) -> List[int]:
    """The product a b over Z, exact, by several NTT primes and the CRT.

    Enough primes are used so that their product exceeds twice the bound
    min(m, n) max|a_i| max|b_j| on the absolute value of a coefficient.
    Cost: O(N log N) operations modulo each of the O(1 + log(bound) / 30) primes.
    """
    if not a or not b:
        return []
    bound = min(len(a), len(b)) * max(abs(x) for x in a) * max(abs(x) for x in b)
    size = 1
    while size < len(a) + len(b) - 1:
        size <<= 1
    count, prod = 0, 1
    for p, _ in NTT_PRIMES:
        if prod > 2 * bound:
            break
        count, prod = count + 1, prod * p
    if prod <= 2 * bound:
        count = (2 * bound).bit_length() // 20 + 1
    primes = _crt_primes(count, size)
    residues, mods, prod = [], [], 1
    for p, g in primes:
        residues.append(multiply_mod(a, b, p, g) + [0] * (len(a) + len(b) - 1))
        mods.append(p)
        prod *= p
    out = []
    for i in range(len(a) + len(b) - 1):
        x, mod = 0, 1
        for r, p in zip(residues, mods):
            # Garner step: x = x + mod * ((r_i - x) / mod mod p)
            t = (r[i] - x) * pow(mod, p - 2, p) % p
            x += mod * t
            mod *= p
        if x > prod // 2:
            x -= prod
        out.append(x)
    return _normalize(out)


def multiply_mod_any(a: Sequence[int], b: Sequence[int], m: int) -> List[int]:
    """The product a b over Z/mZ for any modulus m >= 2, through Z."""
    c = multiply_int([x % m for x in a], [x % m for x in b])
    return _normalize([x % m for x in c])


def multiply_fraction(a: Sequence[Fraction], b: Sequence[Fraction]) -> List[Fraction]:
    """The product a b over Q: clear denominators, multiply over Z, divide back."""
    if not a or not b:
        return []
    da = lcm(*[Fraction(x).denominator for x in a])
    db = lcm(*[Fraction(x).denominator for x in b])
    ia = [int(Fraction(x) * da) for x in a]
    ib = [int(Fraction(x) * db) for x in b]
    c = multiply_int(ia, ib)
    return [Fraction(x, da * db) for x in c]


def multiply_complex(a: Sequence[int], b: Sequence[int]) -> List[int]:
    """The product a b over Z by a floating point FFT over C, rounded.

    Correct only while the rounding error stays below 1/2; for 53-bit doubles
    this holds in practice when N max|a_i| max|b_j| is well below 2^50.
    """
    m, n = len(a), len(b)
    if m == 0 or n == 0:
        return []
    size = 1
    while size < m + n - 1:
        size <<= 1

    def fft(v: List[complex], sign: int) -> None:
        j = 0
        for i in range(1, size):
            bit = size >> 1
            while j & bit:
                j ^= bit
                bit >>= 1
            j |= bit
            if i < j:
                v[i], v[j] = v[j], v[i]
        length = 2
        while length <= size:
            half = length // 2
            for k in range(half):
                w = cmath.exp(sign * 2j * cmath.pi * k / length)
                for s in range(0, size, length):
                    t = w * v[s + k + half]
                    x = v[s + k]
                    v[s + k] = x + t
                    v[s + k + half] = x - t
            length <<= 1

    fa = [complex(x) for x in a] + [0j] * (size - m)
    fb = [complex(x) for x in b] + [0j] * (size - n)
    fft(fa, -1)
    fft(fb, -1)
    fc = [x * y for x, y in zip(fa, fb)]
    fft(fc, 1)
    return _normalize([round(x.real / size) for x in fc[: m + n - 1]])