8.13.4 ACM-ICPC 多项式与生成函数快速数论变换

本文链接：https://blog.csdn.net/tang7mj/article/details/139247126

8.13.4 ACM-ICPC 多项式与生成函数快速数论变换

前置知识：模数论

快速数论变换（Number Theoretic Transform，NTT）是快速傅里叶变换（FFT）在数论基础上的实现。与 FFT 使用复数单位根不同，NTT 使用整数模数下的原根，从而避免了浮点运算误差的问题。

一、引入

与 FFT 类似，NTT 也是用于计算两个多项式乘法的一种高效算法。不同的是，NTT 在整数模数域上进行运算，确保结果是整数，且没有精度损失。

二、数论基础

1. 原根

在模 𝑝p 的整数域上，一个整数 𝑔g 被称为模 𝑝p 的原根，如果 𝑔g 的所有幂模 𝑝p 生成了从 1 到 𝑝−1p−1 的所有整数。形式上，即： 𝑔1,𝑔2,…,𝑔𝑝−1mod 𝑝g1,g2,…,gp−1modp 形成一个完全剩余系。

2. 模数选择

为了应用 NTT，选择一个大质数 𝑝p，使得存在一个整数 𝑔g 是 𝑝p 的原根，且 𝑝=𝑘⋅2𝑚+1p=k⋅2m+1，其中 𝑘k 和 𝑚m 是整数。这确保了 𝑝p 可以支持 2𝑚2m 点的变换。

三、快速数论变换（NTT）

1. 定义

设 𝑎a 是一个长度为 𝑛n 的整数序列，𝑝p 是一个质数，且 𝑝=𝑘⋅2𝑚+1p=k⋅2m+1。定义 𝑔g 为 𝑝p 的原根，NTT 定义为： 𝐴(𝑘)=∑𝑗=0𝑛−1𝑎𝑗⋅𝑔𝑗𝑘mod 𝑝A(k)=∑j=0n−1aj⋅gjkmodp

逆变换为： 𝑎𝑗=1𝑛∑𝑘=0𝑛−1𝐴(𝑘)⋅𝑔−𝑗𝑘mod 𝑝aj=n1∑k=0n−1A(k)⋅g−jkmodp

2. NTT 算法

NTT 的基本步骤与 FFT 类似，包括以下步骤：

点值表示：将多项式转换为点值表示。
乘积计算：计算两个多项式点值表示的乘积。
逆变换：将点值表示转换回系数表示。

3. 代码实现

下面是 NTT 的实现代码：

#include <iostream>
#include <vector>

const int MOD = 998244353;  // 质数模数
const int G = 3;            // 原根

using namespace std;

typedef long long ll;

// 快速幂
ll mod_pow(ll base, ll exp, ll mod) {
    ll result = 1;
    while (exp > 0) {
        if (exp % 2 == 1) {
            result = result * base % mod;
        }
        base = base * base % mod;
        exp /= 2;
    }
    return result;
}

// 递归版 NTT
void ntt(vector<ll> &a, bool invert) {
    int n = a.size();
    if (n == 1) return;

    vector<ll> a0(n / 2), a1(n / 2);
    for (int i = 0; 2 * i < n; i++) {
        a0[i] = a[2 * i];
        a1[i] = a[2 * i + 1];
    }
    ntt(a0, invert);
    ntt(a1, invert);

    ll w = 1;
    ll wn = mod_pow(G, (MOD - 1) / n, MOD);
    if (invert) wn = mod_pow(wn, MOD - 2, MOD);

    for (int i = 0; 2 * i < n; i++) {
        a[i] = (a0[i] + w * a1[i] % MOD) % MOD;
        a[i + n / 2] = (a0[i] - w * a1[i] % MOD + MOD) % MOD;
        if (invert) {
            a[i] = a[i] * mod_pow(n, MOD - 2, MOD) % MOD;
            a[i + n / 2] = a[i + n / 2] * mod_pow(n, MOD - 2, MOD) % MOD;
        }
        w = w * wn % MOD;
    }
}

// 多项式乘法
vector<ll> multiply(vector<ll> const& a, vector<ll> const& b) {
    vector<ll> fa(a.begin(), a.end()), fb(b.begin(), b.end());
    int n = 1;
    while (n < a.size() + b.size()) n <<= 1;
    fa.resize(n);
    fb.resize(n);

    ntt(fa, false);
    ntt(fb, false);
    for (int i = 0; i < n; i++) {
        fa[i] = fa[i] * fb[i] % MOD;
    }
    ntt(fa, true);

    return fa;
}

int main() {
    vector<ll> a = {5, 3, 7};
    vector<ll> b = {7, 2, 1};

    vector<ll> result = multiply(a, b);
    for (ll x : result) {
        cout << x << " ";
    }
    cout << endl;

    return 0;
}