多项式算法——快速数论变换NTT_ntt算法解决什么问题-CSDN博客

本文链接：https://blog.csdn.net/jiahonghao2002/article/details/126038372

多项式算法——快速数论变换NTT

在 FFT 中，我们使用复数计算，未免会出现精度损失，基于原根的快速数论变换 NTT 解决了这个问题。

模意义下的 $n$ 次单位根

在 FFT 中，我们计算出多项式在 $\omega^n = 1$ 的 $n$ 次复数域下的 $n$ 个单位根下的点值来完成 FFT ，这是因为 $n$ 次单位根具有下面的三个性质：

周期性： $n$ 个单位根互不相同，且幂次具有周期性。
消去引理：

$\omega^{dk}_{dn} = e^{\frac{2 dk \pi}{dn}i} = e^{\frac{2 k \pi}{n}i} = \omega^{k}_{n}$
折半引理：

如果 $\gt 0$ 是偶数，那么 $n$ 个 $n$ 次单位复数根的平方的集合，等同于 $n /2$ 个 $n /2$ 次单位复数根的的集合。

根据消去引理 $(\omega_n^k)^2=\omega^k_{\frac{n}{2}}$ ，如果对每个根都进行平方，那么每个不同的数正好出现两次，因为：

$(\omega_n^{k + n/2})^2 = \omega_n^{2k+n} = \omega_n^{2k}\omega_n^n = \omega_n^{2k} = (\omega_n^k)^2 = \omega_{n/2}^k$

我们根据这几条性质，定义模 $p$ （我们只讨论素数模下的NTT）意义下的 $n$ 次主单位根。

设 $g$ 为 $p$ 的一个原根，那么 $g^0,g^1,\ldots,g^{p-1}$ 和 $0,2,\ldots,p-1$ 构成双射关系，也就是 $p$ 的一个简化剩余系，并且周期出现，满足第一条性质。

令 $g_n = g^{\frac{p-1}{n}}$ ，称 $g_n$ 为模 $p$ 意义下的 $n$ 次主单位根。

快速数论变换NTT

有了 $g_n$ 为模 $p$ 意义下的 $n$ 次主单位根，则对于消去引理：

$g^{dk}_{dn} = g^{\frac{(p-1)kd}{dn}} = g^{k}_{n}$

对于折半引理：

$g_n^{k + n/2})^2 = g_n^{2k+n} = g_n^{2k}g_n^n = g_n^{2k} = (g_n^k)^2 = g_{n/2}^k$

在这里 $g_n^n \equiv 1 \mod p$ 可由费马小定理得到。

对比负数域上的单位圆， $\frac{2\pi}{n}$ 为将整个圆等分，而 $\frac{p-1}{n}$ 将剩余系等分。

故我们只需要替换 $w_n$ 为 $g_n$ 就是快速数论变换NTT。

对于逆快速数论变换NTT，也类似逆FTT，在逆FTT中，乘以单位根的倒数，并乘以 $n$ 的倒数。

$a_j = \frac{1}{n}\sum_{k=0}^{n-1}y_k \omega_n^{-kj}$

则对于 NTT ，我们乘以 $g_n^j$ 的逆元和 $n$ 的逆元即可。

要注意的几点为：

NTT 仍存在计算损失，只不过不是浮点误差，当系数超过 $p$ 的时候，那么得到的结果是模完 $p$ 之后的结果。
对于 $p$ 和 $g$ 的选择，我们令 $p$ 尽量是 $2$ 的幂次的倍数加一，可以选择 $p = 998244353, g = 3$ 。

代码

#include <bits/stdc++.h>

using namespace std;

using ll = long long;

#ifdef LLT_DBG
#define FR freopen("in.txt", "r", stdin)
#else
#define FR
#endif

template <ll P>
ll fpow(ll a, ll b)
{
    ll res = 1;
    for (; b; b >>= 1, a = (a * a) % P)
        if (b & 1)
            res = (res * a) % P;
    return res;
}

template <ll G, ll P>
struct NTT
{
    int _n;
    int E;
    vector<int> rev;

    /**
     * @brief 构建一个 NTT 计算器
     *
     * @param n 多项式最高项数
     */
    NTT(int n)
    {
        _n = 1;
        E = 0;
        while (_n < n)
        {
            _n <<= 1;
            E++;
        }
        rev.resize(_n);
        // 逆位置对换
        for (int i = 1; i < _n; i++)
        {
            rev[i] = (rev[i >> 1] >> 1) + ((i & 1) << (E - 1));
        }
    }

    void _rNTT(ll A[], ll k)
    {
        for (int i = 0; i < _n; i++)
            if (i < rev[i])
                swap(A[i], A[rev[i]]);

        for (int e = 1; e <= E; e++)
        {
            int m = 1 << e;

            for (int i = 0; i < _n; i += m)
            {
                int hf = m / 2;
                ll g = 1;
                ll gn = fpow<P>(fpow<P>(G, (P - 1) / m), k);

                for (int j = 0; j < hf; j++)
                {
                    ll x = A[i + j];
                    ll y = (A[i + j + hf] * g) % P;
                    A[i + j] = (x + y) % P;
                    A[i + j + hf] = (x - y) % P;
                    g = (g * gn) % P;
                }
            }
        }
    }

    /**
     * @brief NTT 过程
     *
     * @param A 系数数组
     */
    void doNTT(ll A[])
    {
        _rNTT(A, 1);
    }

    /**
     * @brief NTT 逆过程
     *
     * @param A 点值数组
     */
    void doINTT(ll A[])
    {
        ll ni = fpow<P>(_n, P - 2);
        _rNTT(A, P - 2);
        for (int i = 0; i < _n; i++)
        {
            A[i] = (A[i] * ni) % P;
            A[i] = (A[i] + P) % P;
        }
    }
};

ll A[5000005];
ll B[5000005];

const int mod = 998244353;

void solve()
{
    int n, m;
    cin >> n >> m;

    for (int i = 0; i <= n; i++)
    {
        cin >> A[i];
    }

    for (int i = 0; i <= m; i++)
    {
        cin >> B[i];
    }

    NTT<3, mod> ntt(n + m + 1);

    ntt.doNTT(A);
    ntt.doNTT(B);

    for (int i = 0; i < ntt._n; i++)
    {
        A[i] = (A[i] * B[i]) % mod;
    }

    ntt.doINTT(A);

    for (int i = 0; i < n + m + 1; i++)
    {
        cout << A[i] << " ";
    }
}

int main()
{
    FR;
    solve();
    return 0;
}