【闲话 No.5】 FFT 与 NTT 基础

门的另一端

可爱得很有力量的歌，丘丘人的部分感觉好听炸了。

偶然推开的那一扇门

从何时起变得不再陌生

对面有人 可爱地等

在遇见你的这个世界

时针似乎转得要比平时快些

天空 更清澈了一些

相遇的理由

幻化成了歌

揉进了旅途的景色

为什么此刻

心脏在愉快地蹦着

明明触摸不到的风

却偷偷把汗水吹走

闻不到的饭香

却让我更坚强 更勇敢

明明是受不到的伤

心却还是隐隐作痛鼻腔发痒

明明那是我未曾踏足的彼方

这个地方有一点奇怪

炊烟不会断 花无尽地开

人比猪懒 最爱使唤

偶尔对着天空发个呆

直到云朵透出日落果的色彩

你看 是不是很像

我这个品种

总聚少离多

寂寞能写十本小说

从来没有过

有过这么多

Ye tomo

明明触摸不到的风

却偷偷把汗水吹走

闻不到的饭香

却让我更坚强 更勇敢

明明是受不到的伤

心却还是隐隐作痛鼻腔发痒

想再靠近一些 这熟悉的温暖

没完没了的话

每次听到还是 很喜欢

一首歌太短

在门的另一端 故事还没写完

Yeye dada

Mimi tomo

Mosi mita

Gusha gusha nye

Gusha gusha nye nye

温馨提示

本文更注重讲清 DFT、IDFT 等的实现原理，给出的代码跑得不快且未经良好的封装。

如果您渴望找到一份好用的多项式板子，请移步其它博客。

定义

多项式是一个形如 \(f(x)=\sum_{i}{a_i x^i}\) 的式子，我们也会用它的系数表示它，如表示为 \(<a_0,a_1,\cdots,a_{n-1}>\)。

如无特殊说明，以下 \(n\) 均为多项式的项数。如对于多项式 \(f\)，\(\deg f=n-1\)。

以下默认讨论的是多项式 \(a,b\) 相乘得 \(c\)，其中 \(a,b\) 中长度较大的一个长度为 \(n\)，我们用下标 \(a_i\) 表示 \(i\) 次项的系数，用括号 \(a(x)\) 表示将 \(x\) 代入 \(a\) 求得的值。

我们用 \(\omega_{a}\) 表示 \(a\) 次单位根中辐角最小的一个，也即 \(\cos(\frac{2\pi}{a})+\sin(\frac{2\pi}{a})i\)，\(\omega^{b}_{a}\) 就是它的 \(b\) 次方。容易发现 \(\omega^{a}_{a}=\omega^{0}_{a}=1\)。

引入

当我们要将两个多项式相乘时，即计算 \(c_i=\sum_{j+k=i}{a_j b_k}\)，暴力计算的复杂度为 \(O(n^2)\) 的。考虑如果两个多项式都是点值表示的，那我们容易得到 \(c(x)=a(x)b(x)\)，毕竟对于 \(a(x)b(x)\) 的第 \(i\) 项，应为

\[\begin{aligned}
&\sum_{j+k=i}{a_j x^j b_k x^k}\\=&\sum_{j+k=i}{a_j b_k x^i}\\=&c_i x^i
\end{aligned}
\]

所以可以快速地得到 \(c\) 的点值表示。接下来我们尝试建立点值表示与系数间的关系。

DFT 与 IDFT

我们先来考虑一个性质：

\[\frac{1}{n}\sum_{i=0}^{n-1}\omega_{n}^{ai}=[a \equiv 0 \bmod n]
\]

当 \([a \equiv 0 \bmod n]\) 时，每一个 \(\omega_{n}^{ai}\) 都是 \(1\)，这个式子的值也为 \(1\)。否则，这个式子可以变成等比数列求和的形式，即

\[\begin{aligned}
&\frac{1}{n}\sum_{i=0}^{n-1}\omega_{n}^{ai}\\=&\frac{\omega_{n}^{a(n-1)}\omega_{n}^{a}-\omega_{n}^{0}}{\omega_{n}^{a}-1}\\=&\frac{1-1}{\omega_{n}^{a}-1}\\=&0
\end{aligned}
\]

接下来回到原式，我们考虑

\[\begin{aligned}
&c_i\\=&\sum_{j+k=i}a_j b_k\\=&\sum_{j,k}[j+k=i]a_j a_k
\end{aligned}
\]

我们先考虑另一个和它长得很像的式子：

\[\begin{aligned}
&\sum_{j,k}[j+k \equiv i \bmod n]a_j a_k\\=&\sum_{j,k}[j+k-i \equiv 0 \bmod n]a_j a_k\\=&\sum_{j,k}\frac{1}{n}\sum_{l=0}^{n-1}\omega_{n}^{l(j+k-i)}a_j b_k\\=&\frac{1}{n}\sum_{j,k}\sum_{l=0}^{n-1}\omega_{n}^{-il}\omega_{n}^{jl}a_j\omega_{n}^{kl}k_l\\=&\frac{1}{n}\sum_{l=0}^{n-1}\omega_{n}^{-il}\sum_{j=0}^{n-1}\omega_{n}^{jl}a_j\sum_{k=0}^{n-1}\omega_{n}^{kl}b_k
\end{aligned}
\]

其中第二步代入了上面的结论。

我们发现这个最终的形式非常优美，因为它可以分成几个相似的部分。我们先来尝试建立原式和该式之间的联系。考虑他们的区别是把 \([j+k=i]\) 变成了 \([j+k \equiv i \bmod n]\)，这样做可能使得 \(j+k=i+n\) 的 \(j,k\) 被多算进 \(c_i\) 的贡献。但最终式子 \(c\) 的长度一定是小于 \(2n\) 的，如果我们考虑将两个原式的长度都按 \(2n\) 考虑，就能避免上面的情况，而且其它情况如 \(j+k=i+2n\) 是没有意义的，因为这样 \(j,k\) 必有一个大于等于 \(n\)，此时 \(a_j b_k=0\)。这样，我们就可以通过求下式得到上式的答案。

我们发现 \(d_i=\sum_{j=0}^{n-1}\omega_{n}^{ij}a_j\) 和 \(e_i=\sum_{j=0}^{n-1}\omega_{n}^{ij}b_j\)，就是将 \(\omega_{n}^{i}\) 带入 \(a,b\) 的点值，而 \(c_i=\sum_{j=0}^{n-1}\omega_{n}^{-ij}d_j e_j\) 则建立了一种点值到系数的关系。这两中变换便分别是 DFT 和 IDFT 的一种体现。接下来我们考虑怎么快速计算他们。

FFT

我们先只考虑 DFT，也就是求点值的过程。考虑多项式

\[\begin{aligned}
&f(x)\\=&a_0 +a_1 x+a_2 x^2+a_3 x^3\\=&a_0+a_2 x^2+x(a_1 +a_3 x^2)
\end{aligned}
\]

这样将奇次与偶次分离之后，假设 \(g(x)\) 为 \(<a_0,a_2>\)，\(h(x)\) 为 \(<a_1,a_3>\)，那么

\[f(x)=g(x^2)+x h(x^2)
\]

假设 \(x=\omega_{n}^{i}\)，那么

\[\begin{aligned}
&f(\omega_{n}^{i})\\=&g(\omega_{n}^{2i})+\omega_{n}^{i}h(\omega_{n}^{2i})\\=&g(\omega_{\frac{n}{2}}^{i})+\omega_{n}^{i}h(\omega_{\frac{n}{2}}^{i})
\end{aligned}
\]

假设 \(n\) 为 \(2\) 的平方数，这样我们把长度为 \(n\) 的式子递归为 \(O(\log n)\) 层处理，每层都是若干个长度为 \(a=2^k\) 的多项式，需要分别带入 \(a\) 个值（\(\omega_{a}^{0}\) 到 \(\omega_{a}^{a-1}\)）计算，总复杂度为 \(O(n \log n)\)。

尽管复杂度已经够优秀了，但是这种递归的实现非常不优美。我们假设存储结果的序列为 \(p\)，原序列为 \(a\)，我们先考虑我们是怎么归类系数奇偶的。先把偶数提前，奇数靠后，再将第二个二进制位是 \(0\) 的提前……可以发现这样做相当与将系数按照原位置排序，第 \(i\) 个二进制位为 \(i\) 关键字，可以发现每个数的最终位置就是将它二进制位反转后得到的位置，这样得到的序列在上述排序标准下不存在逆序对。假设 \(i\) 对应的位置为 \(r_i\)，那么我们可以先令 \(p_{r_i}=a_{i}\)，然后我们发现每个上级式子的两个下级式子的系数正好在它的前半部分和后半部分，因为我们对上一关键字排序后对下一关键字的排序是在这个固定的区间内进行的。考虑一开始 \(p_i\) 就是 \(<p_i>\) 代入 \(\omega_{1}^{0}\) 的解，之后每一次运算我们都把这个这个式子代入 \(i\) 次本源单位根的点值存在这个区域的第 \(i\) 个位置，合并到上级式子时发现代入 \(i\) 次本源单位根和 \(i+\frac{n}{2}\) 次本源单位根所需数据的位置是一样的，也就是将要存储这两个数据的位置，因此可以同时处理。最终我们用 \(O(n)\) 的空间倍增解决了这个问题。

对于 IDFT，我们发现 \(\omega_{n}^{-i}=\omega_{n}^{n-i}\)，假设我们像处理 DFT 一样依次代入 \(0,1,\cdots,n-1\) 次本源单位根，发现除了 \(\omega_{n}^{0}\) 代入结果就是 \(c_0\) 以外 \(\omega_{n}^{i}\) 代入结果其实是 \(c_{n-i}\)，所以我们只需按 DFT 处理再除以 \(n\) 并将后 \(n-1\) 项反转即可。

给出一份并不很快的 P3803 【模板】多项式乘法（FFT）通过代码。

代码

#include<bits/stdc++.h>
using namespace std;
const double pi=acos(-1.0);
struct cplx{
    double rl,im;
    friend cplx operator + (cplx x,cplx y){
        return cplx{x.rl+y.rl,x.im+y.im};
    }
    void operator += (cplx x){
        rl+=x.rl,im+=x.im;
    }
    friend cplx operator - (cplx x,cplx y){
        return cplx{x.rl-y.rl,x.im-y.im};
    }
    friend cplx operator * (cplx x,cplx y){
        return cplx{x.rl*y.rl-x.im*y.im,x.rl*y.im+x.im*y.rl};
    }
    void operator *= (cplx x){
        double tem=rl;
        rl=rl*x.rl-im*x.im,im=tem*x.im+x.rl*im;
    }
}w[40];
int r[2200100];
inline void fft(bool opt,int len,cplx* x){
    for(int i=0;i<len;++i) if(r[i]>i) swap(x[i],x[r[i]]);
    for(int l=2,k=1;l<=len;l<<=1,++k){
        cplx tem=cplx{1,0};
        for(int i=0;i<len;++i){
            if(i&(l-1)&(l>>1)){
                i+=(l>>1)-1,tem=cplx{1,0};
                continue;
            }
            cplx tt1=x[i],tt2=tem*x[i|(l>>1)];
            x[i]+=tt2;
            x[i|(l>>1)]=tt1-tt2;
            tem*=w[k];
        }
    }
    if(opt){
        for(int i=0;i<len;++i) x[i].rl/=len;
        reverse(x+1,x+len);
    }
}
struct poly{
    int len;
    cplx data[2200100];
    inline void read(int x){
        for(int i=0;i<x;++i){
            if(data[i].rl>=-0.5&data[i].rl<0) printf("0 ");
            else printf("%.0f ",data[i].rl);
        }
    }
    friend poly operator * (poly x,poly y){
        int tlen=1;
        poly rtr;
        for(;tlen<(max(x.len,y.len)<<1);tlen<<=1);
        for(int i=1;i<tlen;++i){
            r[i]=r[i>>1]>>1;
            if(i&1) r[i]|=(tlen>>1);
        }
        fft(0,tlen,x.data),fft(0,tlen,y.data);
        for(int i=0;i<tlen;++i) rtr.data[i]=x.data[i]*y.data[i];
        fft(1,tlen,rtr.data);
        for(rtr.len=tlen;rtr.data[rtr.len-1].rl==0;--rtr.len);
        return rtr;
    }
}a,b,c;
int main(){
    for(int i=1,j=1;i<24;++i,j<<=1){
        w[i]=cplx{cos(pi/j),sin(pi/j)};
    }
    scanf("%d%d",&a.len,&b.len);
    int tl=a.len+b.len+1;
    ++a.len;
    for(int i=0;i<a.len;++i) scanf("%lf",&a.data[i].rl);
    ++b.len;
    for(int i=0;i<b.len;++i) scanf("%lf",&b.data[i].rl);
    c=a*b,c.read(tl);
    return 0;
}

NTT

考虑计算 \(c_i\) 在模 \(p\) 意义下的结果。我们考虑对于奇素数 \(p\)，如果它有原根 \(g\)，那么

\[\omega_{n} \equiv g^{\frac{\varphi(p)}{n}} \bmod p
\]

其中一个数的原根即为使得 \(b=\varphi(p)\) 是使 \(g^b \equiv 1 \bmod p\) 最小的 \(b\) 的 \(g\)。可以发现 \(g^0,g^1,\cdots,g^{\varphi(p)-1}\) 在模 \(p\) 意义下互不相等，否则对于 \(g^a=g^b\)，\(g^{a-b} \equiv 1 \bmod p\) 与 \(\varphi(p)\) 最小性矛盾。这样我们保证了这 \(n\) 个单位根互不相同，其余按照 FFT 进行即可。

注意 NTT 模数大都形如 \(k2^a+1\)，其中 \(a\) 足够大保证了 \(g^{\frac{\varphi(p)}{n}}\) 运算的正常进行（因为 \(n\) 是 \(2\) 的次方数）。常见的 NTT 模数 \(998244353=7\times 17 \times 2^{23} +1\)，它的一个原根为 \(3\)。

给出一份 P3803 【模板】多项式乘法（FFT）的通过代码。因为保证了各项系数小于 \(10\)，所以可以按模 \(998244353\) 意义下的 NTT 做。

代码

#include<bits/stdc++.h>
using namespace std;
const long long p=998244353;
int r[2200100];
inline long long qpow(long long x,long long y){
    long long rtr=1;
    for(;y;y>>=1){
        if(y&1) rtr=rtr*x%p;
        x=x*x%p;
    }
    return rtr;
}
inline void ntt(bool opt,int len,long long *x){
    for(int i=0;i<len;++i) if(r[i]>i) swap(x[i],x[r[i]]);
    for(int l=2,k=1;l<=len;l<<=1,++k){
        long long tt=qpow(3,(p-1)>>k),tem=1;
        for(int i=0;i<len;++i){
            if(i&(l-1)&(l>>1)){
                i+=(l>>1)-1,tem=1;
                continue;
            }
            long long tt1=x[i],tt2=tem*x[i|(l>>1)]%p;
            x[i]+=tt2; if(x[i]>=p) x[i]-=p;
            x[i|(l>>1)]=tt1-tt2; if(x[i|(l>>1)]<0) x[i|(l>>1)]+=p;
            tem=tem*tt%p;
        }
    }
    if(opt){
        long long tem=qpow(len,p-2);
        for(int i=0;i<len;++i) x[i]=x[i]*tem%p;
        reverse(x+1,x+len);
    }
}
struct poly{
    int len;
    long long data[2200100];
    inline void read(int x){
        for(int i=0;i<x;++i) if(data[i]<0) data[i]+=p;
        for(int i=0;i<x;++i) printf("%lld ",data[i]);
        printf("\n");
    }
    friend poly operator * (poly x,poly y){
        int tlen=1;
        poly rtr;
        for(;tlen<(max(x.len,y.len)<<1);tlen<<=1);
        for(int i=1;i<tlen;++i){
            r[i]=r[i>>1]>>1;
            if(i&1) r[i]+=tlen>>1;
        }
        ntt(0,tlen,x.data),ntt(0,tlen,y.data);
        for(int i=0;i<tlen;++i) rtr.data[i]=x.data[i]*y.data[i]%p;
        ntt(1,tlen,rtr.data);
        for(rtr.len=tlen;!rtr.data[rtr.len-1];--rtr.len);
        return rtr;
    }
}a,b,c;
int main(){
    scanf("%d%d",&a.len,&b.len);
    int tl=a.len+b.len+1;
    ++a.len;
    for(int i=0;i<a.len;++i) scanf("%lld",&a.data[i]);
    ++b.len;
    for(int i=0;i<b.len;++i) scanf("%lld",&b.data[i]);
    c=a*b,c.read(tl);
    return 0;
}

Acknowledgement

感谢 wang54321 掀起机房学习多项式的大潮。

感谢 xrlong 把我带出了学习一大堆根本没有什么关系的前置知识的歧路。

最后感谢傅里叶大神。

Reference

门的另一端 - 百度百科

原根 - OI Wiki

快速傅里叶变换 - OI Wiki

快速数论变换 - OI Wiki

[学习笔记&教程] 信号, 集合, 多项式, 以及各种卷积性变换 (FFT,NTT,FWT,FMT) - rvalue - 博客园

FFT+NTT入门 - CuFeO4 - 博客园