NTT(数论变换)

What · 是什么

数论变换（NTT）是 FFT 在模素数 p 意义下的版本：用模 p 下的本原单位根（原根 g 的幂）替代复数单位根，加减乘全部在整数取模下进行，彻底避免浮点误差。要求 p-1 能被 2^k 整除，常用模数 998244353=119·2^23+1（原根 3），支持长度到 2^23。实现与 FFT 完全同构：位逆序置换+蝶形，只是旋转因子 w_h=g^{(p-1)/h}，逆变换用逆根再乘 n 的逆元。因为结果精确，可处理系数取模的题目、任意模数 NTT(拆系数/多模合并)、以及作为所有多项式算法(逆元/ln/exp/多点求值等)的底层卷积引擎。

Why · 为什么

How · 实现

伪代码

# 迭代 NTT, mod p, 原根 g, n=2^k | p-1
bit_reverse_permute(a)
for h in 2,4,8,...,n:
    w_h = g^((p-1)/h)          # 正变换; 逆变换用逆元
    for i in step h:
        w = 1
        for j in i..i+h/2-1:
            u=a[j]; v=a[j+h/2]*w % p
            a[j]=(u+v)%p; a[j+h/2]=(u-v)%p
            w = w*w_h % p
if invert: 整体乘 n^{-1} mod p

Python 实现

MOD = 998244353
ROOT = 3  # 998244353 = 7*17*2^23 + 1 的原根


def mod_pow(a, b, mod=MOD):
    r = 1
    while b:
        if b & 1:
            r = r * a % mod
        a = a * a % mod
        b >>= 1
    return r


def ntt(a, invert=False):
    """NTT(原根实现, 蝴蝶操作)。a: list[int] mod MOD, 长度须为 2 的幂。"""
    n = len(a)
    a = a[:]
    # 位逆序置换
    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]
    h = 2
    while h <= n:
        # w_h = g^((MOD-1)/h) 或逆元
        w_h = mod_pow(ROOT, (MOD - 1) // h)
        if invert:
            w_h = mod_pow(w_h, MOD - 2)
        for i in range(0, n, h):
            w = 1
            for j in range(i, i + h // 2):
                u = a[j]
                v = a[j + h // 2] * w % MOD
                a[j] = (u + v) % MOD
                a[j + h // 2] = (u - v) % MOD
                w = w * w_h % MOD
        h <<= 1
    if invert:
        inv_n = mod_pow(n, MOD - 2)
        for i in range(n):
            a[i] = a[i] * inv_n % MOD
    return a


def poly_mul_ntt(a, b, mod=MOD):
    n = 1
    need = len(a) + len(b) - 1
    while n < need:
        n <<= 1
    fa = a[:] + [0] * (n - len(a))
    fb = b[:] + [0] * (n - len(b))
    fa = ntt(fa)
    fb = ntt(fb)
    for i in range(n):
        fa[i] = fa[i] * fb[i] % mod
    fa = ntt(fa, True)
    return fa[:need]


def poly_mul_naive(a, b, mod=MOD):
    res = [0] * (len(a) + len(b) - 1)
    for i, x in enumerate(a):
        for j, y in enumerate(b):
            res[i + j] = (res[i + j] + x * y) % mod
    return res[:len(a) + len(b) - 1]


if __name__ == "__main__":
    import random
    # 与朴素卷积对拍
    for _ in range(50):
        la = random.randint(1, 10)
        lb = random.randint(1, 10)
        a = [random.randint(0, MOD - 1) for _ in range(la)]
        b = [random.randint(0, MOD - 1) for _ in range(lb)]
        assert poly_mul_ntt(a, b) == poly_mul_naive(a, b), (a, b)
    # OI-wiki 引入样例
    assert poly_mul_ntt([7, 3, 5], [1, 2, 7]) == [7, 17, 60, 31, 35]
    # 单位根性质: g_n^(n/2) = -1
    g4 = mod_pow(ROOT, (MOD - 1) // 4)
    assert g4 * g4 % MOD == MOD - 1
    print("NTT tests passed")

易错点

• 模数不是 NTT 友好素数(2^k∤p-1)时直接套用会错
• 长度超过模数支持的 2^k 上界(如 998244353 只支持 2^23)
• 逆变换忘记乘 n 的逆元
• 旋转因子方向与 FFT 混淆(正变换与逆变换用反)
• 任意模数题目忘记拆系数/CRT, 用普通 NTT 出错误答案

例题

信源

• OI Wiki✓ https://oi-wiki.org/math/poly/ntt/
• OI Wiki✓ https://raw.githubusercontent.com/OI-wiki/OI-wiki/master/docs/math/poly/ntt.md
• CP-Algorithms✓ https://cp-algorithms.com/algebra/fft.html