多项式（2）— 快速数论变换（NTT）

前言

前置知识：

• 阶、原根、离散对数;
• 多项式（1）— 快速傅里叶变换（FFT）。

FFT 已经能很好的在 $O(n\log n)$ 内解决卷积问题，但是 FFT 用到的单位根和三角函数有关，只能采用浮点数储存，必定会带来精度误差。

我们尝试不在复数域内，而是在模剩余系内寻求一种新的单位根来解决问题，这样带来了两个好处：

1. 不再有误差（只要大小没有超出模数）；
2. 如果题目要求取模，一并可以完成取模的任务。

新型单位根

我们回看 FFT，FFT 中单位根 $w_n^k$ 满足了如下性质：

1. $w_n^k = \left(w_n^1\right)^k$；
2. $w_n^0,w_n^1,\ldots,w_n^{n-1}$ 互不相同；
3. $w_n^{k+n/2}=-w_n^k$；
   这个性质可以导出 $w_n^{k\bmod n}=w_n^k$。
4. $w_{2n}^{2k}=w_n^k$。

只要满足了以上性质，就可以作为 FFT 中的单位根。

假设我们的模数是质数 $p$，其一个原根为 $g$。

定义 $w_n^1 = g^{(p-1)/n}$，$w_n^k = \left(w_n^1\right)^k$，可以证明其满足如上所有性质。（保证 $n$ 是 $p-1$ 的因数）

证明：

1. 由我们的定义得证。
2. 由阶的性质可知 $w_n^1$ 的阶为 $\dfrac{p-1}{\gcd(p-1,(p-1)/n)} = \dfrac{p-1}{(p-1)/n} = n$，则得证。
3. 只需证明 $x_n^{n/2}\equiv -1 \pmod p$ 即可。
   考虑引理（二次探测定理）：如果 $p$ 是素数，$x^2\equiv 1\pmod p$ 的解为 $x_1=1,x_2=p-1$。
   已知 $\left(x_n^{n/2}\right)^2 \equiv 1 \pmod p$，显然其不能等于 $1$（否则与阶的性质矛盾），所以只能等于 $p-1$。
4. $\text{左边}=g^{2k(p-1)/2n}=g^{k(p-1)/n}=\text{右边}$。

由此，我们就成功消灭了复数单位根。我们把采用这种单位根用来卷积的算法叫做快速数论变换（NTT）。

NTT 实现

我们考虑模数 $p$ 的选取。
由上文我们知道 $p-1$ 一定要是我们分治长度 $n$ 的倍数，我们希望 $p-1$ 中含有尽可能多的 $2$ 的因子。

一个常用模数 $998244353 = 119 \times 2^{23}+1$，这就表明，用 $998244353$ 做模数的 NTT 的序列长度最长为 $2^{23}$。（当然不会这么长的，要不然先 TLE 了）

其一个原根为 $3$。

NTT 相关质数与原根表 (Miskcoo’s Space)

我们用 $998244353$ 作为模数实现 NTT，模板代码如下（细节见注释）：

#include<bits/stdc++.h>
using namespace std;
using ll = long long;

const int MOD = 998244353;
int ksm(int a,int b){ // 快速幂
    int r = 1;
    while(b){
        if(b&1)r=(ll)r*a%MOD;
        a=(ll)a*a%MOD;
        b>>=1;
    }
    return r;
}
const int G = 3, INVG = ksm(G,MOD-2); // 原根和原根的逆

const int N = 4e6+5;
int rev[N];
void NTT(int* f,int n,bool flag){
    for(int i=0;i<n;i++){
        if(i<rev[i])swap(f[i],f[rev[i]]);
    }
    for(int p=2;p<=n;p<<=1){
        int g = ksm(flag?INVG:G, (MOD-1)/p); // 换成新的单位根
        for(int j=0;j<n;j+=p){
            int gnow = 1;
            for(int k=j;k<(j|p>>1);k++){
                int t = (ll)f[k|p>>1] * gnow % MOD;
                f[k|p>>1] = f[k] - t; if(f[k|p>>1]<0)f[k|p>>1]+=MOD; // 注意取模
                f[k] = f[k] + t;      if(f[k]>=MOD)f[k]-=MOD;
                gnow = (ll)gnow * g % MOD;
            }
        }
    }

    int invN = ksm(n, MOD-2);
    if(flag)for(int i=0;i<n;i++)f[i]=(ll)f[i]*invN%MOD; // 改成逆元
}

int f[N],g[N]; // 主函数完全没有变动
int main(){
    ios::sync_with_stdio(0);cin.tie(0);

    int n,m;cin>>n>>m;
    for(int i=0;i<=n;i++)cin>>f[i];
    for(int i=0;i<=m;i++)cin>>g[i];

    for(m+=n,n=1;n<=m;n<<=1);
    for(int i=0;i<n;i++){
        rev[i]=rev[i>>1]>>1;
        if(i&1)rev[i]|=n>>1;
    }

    NTT(f,n,0);NTT(g,n,0);
    for(int i=0;i<n;i++)f[i]=(ll)f[i]*g[i]%MOD;
    NTT(f,n,1);
    for(int i=0;i<=m;i++){
        cout<<f[i]<<' ';
    }
    cout<<'\n';
    return 0;
}

例题（利用 NTT 解决计数问题）

ABC390G - Permutation Concatenation

题意

给定正整数 $n$，求 $1\sim n$ 的 $n!$ 个排列拼接起来得到的数字和，对 $998244353$ 取模。
比如 $n=3$，答案为 $123+132+213+231+312+321=1332$。

解法

考虑拆开考虑每一个数的贡献。
设 $c_i$ 表示 $1\sim n$ 中长度为 $i$ 的数个数，$l_i$ 表示 $i$ 的长度。

假设当前数为 $m$，定义 $d_i$ 表示除了 $m$ 其他数中长度为 $i$ 的数个数。

枚举后面有 $j$ 个数，前面就有 $n-1-j$ 个数，贡献即为

$$\sum_{j=0}^{n-1}j!(n-j-1)!\sum_{i_1 + i_2+\ldots + i_6 = j}^{i_k<d_k}\prod{\binom{d_k}{i_k}}10^{\sum_{k=1}^6{ki_k}}$$

两个阶乘表示了前后的数的顺序排序，组合数表示了选法，$10$ 的幂表示了 $m$ 后面有几位。

我们尝试处理出阶乘后面的一项，也就是后面有 $j$ 个数的取法带上 $10$ 的幂的贡献。

定义形式幂级数

$$\begin{aligned} & F_i=\sum_{j=0}^{c_i}10^{ij}\binom{c_i}{j}x^j \\ & G_i=\sum_{j=0}^{c_i-1}10^{ij}\binom{c_i-1}{j}x^j \\ & H_t=\frac{G_t\prod_{i=1}^6F_i}{F_t} \\ \end{aligned}$$

这样 $[x^j]H_i$ 表示当前数长度为 $i$，后面有 $j$ 个数的贡献，则最后的系数可以写作：

$$\sum_{j=0}^nj!(n-j-1)![x^j]H_t$$

在实现方面，要计算 $H_1\sim H_6$，直接把 $F_1\sim F_6, G_1\sim G_6$ 都转化成点值形式，然后 $H$ 的点值直接乘出来，最后统一化回系数形式即可。

这样只需要 $18$ 次 NTT，复杂度 $O(n\log^2 n)$，通过了本题。

代码

#include<bits/stdc++.h>
using namespace std;

template <class T>
using must_int = enable_if_t<is_integral<T>::value, int>;
template <unsigned umod>
struct modint {
    static constexpr int mod = umod;
    unsigned v;
    modint() : v(0) {}
template <class T, must_int<T> = 0>
     modint(T _v) {int x = _v % (int)umod; v = x < 0 ? x + umod : x;}
     modint operator+() const { return *this; }
     modint operator-() const { return modint() - *this; }
     friend int raw(const modint &self) { return self.v; }
     friend ostream &operator<<(ostream &os, const modint &self) { return os << raw(self);}
     modint &operator+=(const modint &rhs) {v += rhs.v;if (v >= umod) v -= umod;return *this;}
     modint &operator-=(const modint &rhs) {v -= rhs.v;if (v >= umod) v += umod;return *this;}
     modint &operator*=(const modint &rhs) {v = 1ull * v * rhs.v % umod; return *this;}
     modint &operator/=(const modint &rhs) { return *this *= rhs.inv(); }
     modint inv() const {
         assert(v);
         static unsigned lim = 1 << 21;
         static vector<modint> inv{0, 1};
         if (v >= lim) return qpow(*this, mod - 2);
         inv.reserve(v + 1);
         while (v >= inv.size()) {
             int m = inv.size();
             inv.resize(m << 1);
             for (int i = m; i < m << 1; i++)inv[i] = (mod - mod / i) * inv[mod % i];
         }
         return inv[v];
     }
    template <class T, must_int<T> = 0>
    friend modint qpow(modint a, T b) {modint r = 1;for (; b; b >>= 1, a *= a)if (b & 1) r *= a;return r;}
    friend modint operator+(modint lhs, const modint &rhs) { return lhs += rhs; }
    friend modint operator-(modint lhs, const modint &rhs) { return lhs -= rhs; }
    friend modint operator*(modint lhs, const modint &rhs) { return lhs *= rhs; }
    friend modint operator/(modint lhs, const modint &rhs) { return lhs /= rhs; }
    friend bool operator==(const modint &lhs, const modint &rhs) { return lhs.v == rhs.v; }
    friend bool operator!=(const modint &lhs, const modint &rhs) { return lhs.v != rhs.v; }
};
const int MOD = 998244353;
using mint = modint<MOD>;

inline mint C(int n,int r){
    static vector<mint> f{1, 1}, finv{1, 1};
    while(n >= f.size()){
        int m = f.size();
        f.resize(m << 1), finv.resize(m << 1);
        for (int i = m; i < m << 1; i++)f[i]=f[i-1]*i,finv[i]=finv[i-1]/i;
    }
    return (r<0)?0:(f[n]*finv[r]*finv[n-r]);
}

const mint G = 3, INVG = mint{3}.inv();

const int N = 2e6+5;
mint pow10[N];
void NTT(vector<mint>& f,int n,bool flag){
    vector<int> rev(n);
    for(int i=0;i<n;i++){
        rev[i]=rev[i>>1]>>1;
        if(i&1)rev[i]|=n>>1;
    }
    for(int i=0;i<n;i++){
        if(i<rev[i])swap(f[i],f[rev[i]]);
    }
    for(int p=2;p<=n;p<<=1){
        mint g = qpow((flag?INVG:G),(MOD-1)/p); 
        for(int j=0;j<n;j+=p){
            mint gnow = 1;
            for(int k=j;k<(j|p>>1);k++){
                mint t = f[k|p>>1] * gnow;
                f[k|p>>1] = f[k] - t; 
                f[k] = f[k] + t;      
                gnow *= g;
            }
        }
    }

    mint invN = qpow(mint{n}, MOD-2);
    if(flag)for(int i=0;i<n;i++)f[i]*=invN;
}

vector<mint> f[7],g[7],h[7];
int cnt[7];

int len[N];
mint fact[N];
mint coef[7];

int main(){
    ios::sync_with_stdio(0);cin.tie(0);
    int n;cin>>n;
    for(int i=1;i<=n;i++){
        len[i]=len[i/10]+1;
        cnt[len[i]]++;
    }
    pow10[0]=1;
    for(int i=1;i<N;i++)pow10[i]=pow10[i-1]*10;
    fact[0]=1;
    for (int i =1; i <=n; i++)fact[i]=fact[i-1]*i;

    int m;
    for(m=n,n=1;n<=m;n<<=1);


    for(int i=1;i<=6;i++){
        f[i].resize(n); g[i].resize(n);
        for(int j=0;j<=cnt[i];j++)  f[i][j]=C(cnt[i],j)*pow10[i*j];
        for(int j=0;j<=cnt[i]-1;j++)g[i][j]=C(cnt[i]-1,j)*pow10[i*j];
    
        NTT(f[i],n,0);NTT(g[i],n,0);
    }


    for(int i=1;i<=6;i++){
        h[i].resize(n);
        for(int k=0;k<n;k++){
            h[i][k]=1;
            for(int j=1;j<i;j++)h[i][k]*=f[j][k];
            for(int j=i+1;j<=6;j++)h[i][k]*=f[j][k];
            h[i][k]*=g[i][k];
        }

        NTT(h[i],n,1);

        for(int j=0;j<m;j++){
            coef[i] += fact[j]*fact[m-1-j]*h[i][j];
        }
    }

    mint ans = 0;
    for(int i=1;i<=n;i++){
        ans += coef[len[i]] * i;
    }

    cout << ans << '\n';
    return 0;
}