fft-polynomial-multiplication

Answer

import math
from typing import List, Union

def is_power_of_two(n: int) -> bool:
    return n > 0 and (n & (n - 1)) == 0

def next_power_of_two(n: int) -> int:
    return 1 << (n - 1).bit_length()

def fft(a: List[complex], invert: bool = False) -> List[complex]:
    """
    Cooley-Tukey radix-2 FFT (decimation-in-time).
    - a: list of complex numbers, length must be a power of two.
    - invert: if True, compute the inverse FFT (÷n scaling applied).
    """
    n = len(a)
    if n <= 1:
        return a[:]

    # Bit-reversal permutation
    j = 0
    for i in range(1, n):
        bit = n >> 1
        while j & bit:
            j ^= bit
            bit >>= 1
        j ^= bit
        if i < j:
            a[i], a[j] = a[j], a[i]

    # Iterative radix-2 DIT butterfly stages
    length = 2
    while length <= n:
        ang = 2 * math.pi / length * (-1 if invert else 1)
        wlen = complex(math.cos(ang), math.sin(ang))

        for i in range(0, n, length):
            w = 1 + 0j
            half = length // 2
            for j in range(i, i + half):
                u = a[j]
                v = a[j + half] * w
                a[j] = u + v
                a[j + half] = u - v
                w *= wlen
        length <<= 1

    if invert:
        for i in range(n):
            a[i] /= n

    return a

def multiply_polynomials(p: List[Union[float, int]],
                         q: List[Union[float, int]]) -> List[float]:
    """
    Multiply two polynomials via FFT-based convolution.
    Zero-pads to power-of-two length, forward FFTs, pointwise multiply,
    inverse FFT, rounds to integer coefficients.
    """
    n = next_power_of_two(max(len(p), len(q)) * 2)

    # Zero-pad
    fa = [complex(x, 0) for x in p] + [0j] * (n - len(p))
    fb = [complex(x, 0) for x in q] + [0j] * (n - len(q))

    # Forward FFTs
    fft(fa, invert=False)
    fft(fb, invert=False)

    # Pointwise multiply in frequency domain
    for i in range(n):
        fa[i] *= fb[i]

    # Inverse FFT
    fft(fa, invert=True)

    # Round to real integer coefficients
    deg = len(p) + len(q) - 2
    return [round(fa[i].real) for i in range(deg + 1)]

How it works: - Bit-reversal permutation reorders inputs for in-place DIT computation. - Iterative butterfly stages combine pairs using twiddle factors ( e^{\pm 2\pi i k / L} ), with sign flipped for inverse. - Polynomial multiplication: zero-pad both coefficient lists to a power of two (at least deg(p)+deg(q)+1), forward FFT each, pointwise multiply, inverse FFT, and round real components.