NTT 入门

一只绝帆 · 2025-01-24 17:13:55 · 个人记录

牢记目标：对于一个给定的 m-1 次的多项式（m=2^n），我们要求出 m 个点 \omega_{m}^{k} 代进去的点值，有了点值我们可以进行多项式乘法等内容。

把多项式奇偶分类，F(x)=F_0(x^2)+xF_1(x^2)。

由于我们递归下去要规模减半，考虑对点值分成两半，显然 \omega_{m}^{k+(m/2)}=-\omega_{m}^k，所以偶次多项式的后一半点值等价于前一半点值，奇次多项式的后一半点值是前一半点值取反。

把点值还原成多项式（\rm IDFT）时，结论是指数上多一个负号，得到的结果应除以长度。

考虑构造一个多项式 G(x)=\sum_{i=0}^{m-1}F(\omega_{m}^i)x^i，然后对其推式子。

这里实质是一个单位根反演，我们发现 j=k 时每个 i 的贡献均为 1，j\ne k 时所有 i 的贡献和为 0，所以 G(\omega^{-k})=ma_k，所以求解 \rm IDFT 时只需把 \omega 换成 \omega^{-1} 并在最后除掉总长度即可。

值得注意的是，如果我们并不把指数取反，而是代入原来的指数，那根据单位根的性质，其实我们相当于倒着代入了负指数，所以最后加一个 reverse 即可。

注意是从 1 开始的 \rm reverse，因为 \omega^0=\omega^{-0}。

单位根换成原根后我们可以得到一个递归写法：

const int g=3,p=998244353;
ll ksm(ll a,ll n=p-2,ll res=1) {
    for(;n;n>>=1,a=a*a%p) if(n&1) res=res*a%p;
    return res;
}
void NTT(ll w,vec<ll> &a) {
    if(sz(a)==1) return;
    int n=sz(a);
    vec<ll> b[2]{vec<ll>(n/2),vec<ll>(n/2)};
    F(i,0,n-1) b[i&1][i>>1]=a[i];
    NTT(w*w%p,b[0]);NTT(w*w%p,b[1]);
    ll r=1;
    F(i,0,n/2-1) {
        a[i]=(b[0][i]+r*b[1][i])%p;
        a[i+(n/2)]=(b[0][i]+(p-r)*b[1][i])%p;
        (r*=w)%=p;
    } 
}
void DFT(vec<ll> &a) {NTT(ksm(g,(p-1)/sz(a)),a);}
void IDFT(vec<ll> &a) {
    NTT(ksm(ksm(g,(p-1)/sz(a))),a);
    ll iv=ksm(sz(a));
    for(ll &x:a) (x*=iv)%=p;
}

为了减小常数，我们通常要将分治的过程循环化，我们发现我们先从上往下进行了奇偶分类，再自底向上合并。

如果我们要用循环代替递归，此时我们如何处理“奇偶分类”？

我们发现，奇偶分类恰恰是把最低位提到了最高位，所有奇偶分类都做完之后相当于二进制 \rm reverse。

这个过程叫蝴蝶变换，我们可以用 \frac i 2 的 \rm reverse 推出 i 的 \rm reverse，由此我们可以求出蝴蝶变换的位置：

F(i,0,(1<<n)-1) r[i]=r[i>>1]>>1|((i&1)<<n-1);

处理完奇偶分类后，我们就从底向上进行合并过程，这部分就直接抄过来即可。

void NTT(vec<ll> &a,bool type) {
    int n=sz(a),l=__lg(n);
    vec<int> r(n);
    F(i,0,n-1) r[i]=r[i>>1]>>1|(i&1)<<l-1;
    F(i,0,n-1) if(i<r[i]) swap(a[i],a[r[i]]);
    for(int m=1;m<n;m*=2) {
        ll w=ksm(type?g:ksm(g),(p-1)/(m*2));
        for(int j=0;j<n;j+=m*2) {
            ll r=1;
            F(i,j,j+m-1) {
                ll x=a[i],y=r*a[i+m]%p;
                a[i]=(x+y)%p;a[i+m]=(x+p-y)%p;
                (r*=w)%=p;
            }
        }
    }
    if(!type) {
        ll iv=ksm(n);
        F(i,0,n-1) (a[i]*=iv)%=p;
    }
}
void DFT(vec<ll> &a) {NTT(a,1);}
void IDFT(vec<ll> &a) {NTT(a,0);}

这是一份符合大多数人刻板印象的 NTT 模板。

DIT 与 DIF

上文所述的 \rm DFT 方式属于 \rm DIT，按时域抽取，这里的按时域抽取指的是把多项式系数进行奇偶分类。

事实上，我们存在另一种 \rm DFT 方式：\rm DIF，按频域抽取，指的是把代入的点进行奇偶分类。

考虑 \omega 的奇偶性，也就是把多项式按照前后劈开：

当 2\mid k 时：