再探快速傅里叶变换(FFT)学习笔记(其三)(循环卷积的Bluestein算法+分治FFT+FFT的优化+任意模数NTT)

再探快速傅里叶变换(FFT)学习笔记(其三)(循环卷积的Bluestein算法+分治FFT+FFT的优化+任意模数NTT)

写在前面

为了不使篇幅过长，预计将把学习笔记分为四部分:

1. DFT,IDFT,FFT的定义,实现与证明:快速傅里叶变换(FFT)学习笔记(其一)
2. NTT的实现与证明:快速傅里叶变换(FFT)学习笔记(其二)
3. 任意模数NTT与FFT的优化技巧
4. 多项式相关操作

一些约定

1. \([p(x)]=\begin{cases}1,p(x)为真 \\ 0,p(x)为假 \end{cases}\)
2. 本文中序列的下标从0开始
3. 若\(s\)是一个序列,\(|s|\)表示\(s\)的长度
4. 若大写字母如\(F(x)\)表示一个多项式,那么对应的小写字母如\(f\)表示多项式的每一项系数，即\(F(x)=\sum_{i=0}^{n-1} f_ix^i\)

循环卷积

DFT卷积的本质

考虑在(其一)中提到的卷积的定义式。

\[c_{r}=\sum_{p, q}[(p+q) \bmod n=r] a_{p} b_{q} \tag{1.1}\]

我们一般做FFT时忽略了式子中的\(\bmod\),其实它是在\(\bmod 2^q\)的意义下的循环卷积,只是因为\(|a|,|b|,|c|<2^q\),所以取不取模都没什么影响。

如果序列长度\(n\)是2的整数次幂,那么直接做就可以了。

如果序列长度\(n\)不是2的整数次幂考虑暴力的做法:先做一次普通FFT,再把\(c_{k+n}\)加到\(c_k\)上。但是这样显然不是很优秀。下面给出了一种在\(O(n \log n)\)的时间内实现任意长度循环卷积的算法:Bluestein’s Algorithm

Bluestein’s Algorithm

注:原论文的推导可能有误

考虑DFT的式子

\[\begin{aligned} a'_i&=\sum_{j=0}^{n-1} a_j \omega_n^{ij} \\&=\sum_{j=0}^{n-1} a_j \omega_n^{\frac{-(i-j)^2+i^2+j^2}{2}} \\&= \omega_n^{\frac{i^2}{2}} \sum_{j=0}^{n-1}a_j \omega_n^{\frac{j^2}{2}} \omega_n^{-\frac{(i-j)^2}{2}}\end{aligned}\]

不妨设

\[x_j=a_j \omega_n^{\frac{j^2}{2}}=a_j(\cos\frac{j^2\pi}{n}+ \text{i}\sin{\frac{j^2\pi}{n}})\]

$y_j=\omega_n^{-\frac{j^2}{2}}= \cos \frac{\pi j^2}{n}-\text{i}\sin \frac{\pi j^2}{n} $

那么\(a_i'=\omega_n^{\frac{j^2}{2}}\sum_{j=0}^{n-1} x_j y_{i-j}\)

这已经很类似卷积的形式了,但是注意到\(j\)的上界是\(n-1\)而不是\(i\),\(j-i\)可能为负数。那么我们把\(y\)数组的长度扩大到\(2n\),定义:

$y_j=\omega_n^{-\frac{(j-n)^2}{2}}= \cos \frac{\pi (j-n)^2}{n}-\text{i}\sin \frac{\pi (j-n)^2}{n} $.

这样\(j<n\)的时候就对应了\(j-i\)为负数的情形,\(j\geq n\)就对应了\(j-i\)为正的情形。然后对\(x\)和\(y\)用一般的FFT,最后的答案存储在\(i+n\)的位置上，也就是说真正的\(a'_i\)实际上对应了乘积结果的\((x \cdot y)_{i+n}\)

这样，我们就只做了一次FFT就求出了任意长度循环卷积。逆变换同理,只是换成共轭复数。注意到在上述的推导中我们没有用到单位根\(\omega\)的任何性质,因此这里的\(\omega\)可以换成任意复数\(z\),这样的变换称为Chirp Z-Transform,CZT.可见,CZT实际上是DFT的广义形式。

代码实现见下方模板题

例题

这是Bluestein算法的模板题

[POJ 2821] 给出两个长度为\(n\)的序列\(B,C\),已知\(A\)和\(B\)的循环卷积为\(C\),求\(A\).

\(n<2^{17}\)

代码:

#include<iostream>
#include<cstdio>
#include<cstring>
#include<cmath>
#define maxn (1<<17)
const double pi=acos(-1.0);
using namespace std; 
struct com{
    double real;
    double imag;
    com(){
        
    } 
    com(double _real,double _imag){
        real=_real;
        imag=_imag;
    }
    com(double x){
        real=x;
        imag=0;
    }
    void operator = (const com x){
        this->real=x.real;
        this->imag=x.imag;
    }
    void operator = (const double x){
        this->real=x;
        this->imag=0;
    }
    friend com operator + (com p,com q){
        return com(p.real+q.real,p.imag+q.imag);
    }
    friend com operator + (com p,double q){
        return com(p.real+q,p.imag);
    }
    void operator += (com q){
        *this=*this+q;
    }
    void operator += (double q){
        *this=*this+q;
    }
    friend com operator - (com p,com q){
        return com(p.real-q.real,p.imag-q.imag);
    }
    friend com operator - (com p,double q){
        return com(p.real-q,p.imag);
    }
    void operator -= (com q){
        *this=*this-q;
    }
    void operator -= (double q){
        *this=*this-q;
    }
    friend com operator * (com p,com q){
        return com(p.real*q.real-p.imag*q.imag,p.real*q.imag+p.imag*q.real);
    }
    friend com operator * (com p,double q){
        return com(p.real*q,p.imag*q);
    } 
    void operator *= (com q){
        *this=(*this)*q;
    }
    void operator *= (double q){
        *this=(*this)*q;
    }
    friend com operator / (com p,double q){
        return com(p.real/q,p.imag/q);
    } 
    void operator /= (double q){
        *this=(*this)/q;
    } 
    friend com operator / (com p,com q){//复数的除法,类似解二元一次方程,代入复数乘法公式解出答案
        return com((p.real*q.real+p.imag*q.imag)/(q.real*q.real+q.imag*q.imag),(p.imag*q.real-p.real*q.imag)/(q.real*q.real+q.imag*q.imag));
    }
    void print(){
        printf("%lf + %lf i ",real,imag);
    }
};


void fft(com *x,int *rev,int n,int type){
    for(int i=0;i<n;i++) if(i<rev[i]) swap(x[i],x[rev[i]]);
    for(int len=1;len<n;len*=2){
        int sz=len*2;
        com wn1=com(cos(2*pi/sz),type*sin(2*pi/sz));
        for(int l=0;l<n;l+=sz){
            int r=l+len-1;
            com wnk=1;
            for(int i=l;i<=r;i++){
                com tmp=x[i+len];
                x[i+len]=x[i]-wnk*tmp;
                x[i]=x[i]+wnk*tmp;
                wnk=wnk*wn1;
            }
        }
    }
    if(type==-1) for(int i=0;i<n;i++) x[i]/=n;
} 
void bluestein(com *a,int n,int type){ 
    static com x[maxn*4+5],y[maxn*4+5];
    static int rev[maxn*4+5];
    memset(x,0,sizeof(x));
    memset(y,0,sizeof(y));
    int N=1,L=0;
    while(N<n*4){
        L++;
        N*=2;
    }
    for(int i=0;i<N;i++) rev[i]=(rev[i>>1]>>1)|((i&1)<<(L-1));
    for(int i=0;i<n;i++) x[i]=com(cos(pi*i*i/n),type*sin(pi*i*i/n))*a[i];
    for(int i=0;i<n*2;i++) y[i]=com(cos(pi*(i-n)*(i-n)/n),-type*sin(pi*(i-n)*(i-n)/n));
    fft(x,rev,N,1);
    fft(y,rev,N,1);
    for(int i=0;i<N;i++) x[i]*=y[i];
    fft(x,rev,N,-1);
    for(int i=0;i<n;i++){
        a[i]=x[i+n]*com(cos(pi*i*i/n),type*sin(pi*i*i/n));
        if(type==-1) a[i]/=n;//一定记得除以n,因为做一次Bluestein相当于一次FFT,IFFT最后要除n,这里也要除n 
    } 
}
void div(com *a,com *b,com *c,int n){//求解A*B=C 
    bluestein(b,n,1);
    bluestein(c,n,1);
    for(int i=0;i<n;i++) a[i]=c[i]/b[i];
    bluestein(a,n,-1);
}

int n;
com a[maxn+5],b[maxn+5],c[maxn+5];
int main(){
    scanf("%d",&n);
    for(int i=0;i<n;i++) scanf("%lf",&b[i].real);
    for(int i=0;i<n;i++) scanf("%lf",&c[i].real);
    div(a,b,c,n);
    for(int i=0;i<n;i++) printf("%.4f\n",a[i].real);
}

分治FFT

//填坑中

FFT的弱常数优化

下面介绍一些优化FFT的常数的技巧。虽然这些技巧都只是对FFT的一些小优化，但是在某些题目中优化效果极其明显。

复杂算式中减少FFT次数

如果我们要计算一个复杂的多项式,如\(A(x)=B(x)C(x)+D(x)E(x)\)

最简单的方法是分别计算\(B(x)C(x)\)和\(D(x)E(x)\),这样需要做6次FFT. 但是如果先对\(B,C,D,E\)做DFT,然后直接用点值表达式计算\(a_i=b_ic_i+d_ie_i\),再把\(a\)IDFT回去。这样只需要做5次FFT,且多项式越复杂，这样的常数就越优秀。

例题

[BZOJ 3771] Triple(FFT+容斥原理+生成函数)

利用循环卷积

考虑对于两个长度为\(n\)的序列\(a,b\),计算它们的卷积\(c\)的第\(0.5n\)项到第\(1.5n\)项。传统的方法是补0扩充到\(2n\)的序列。但是因为FFT求得实际上是我们已经提到过的循环卷积,所以如果只补0到\(1.5n\)(上取整),对第\(0.5n\)项到第\(1.5n\)项无影响

在基于牛顿迭代的算法中，能起到较明显的优化作用。会在(其四)中详细介绍这些算法。

小范围暴力

由于FFT的常数较大。在数据范围较小的时候甚至不如\(O(n^2)\)的暴力卷积的优秀。因此在做多次FFT和分治FFT的时候,如果当前的序列长度较小，可以采用暴力算法。

例题

[BZOJ 3509] [CodeChef] COUNTARI (FFT+分块)

快速幂乘法次数的优化

这个东西实际上比较鸡肋。因为多项式快速幂可以通过多项式\(\ln\)和\(\exp\)优化到\(O(n \log n)\).但是为了应对考场上时间不够的情况,我们来考虑如何通过简单的实现来减少\(O(n \log^2n)\)的倍增快速幂的复杂度。

倍增法的思路是根据前面算过的乘积快速算出当前的乘积,如\(1 \to 2 \to 4 \to 8\).最坏情况下需要\(2 \log_2n+C\)次乘法。但这并不是下界。我们定义additional chain为一条链,最开始是1,后一个数减前一个数的差是链上这个是前面的某一个数。例如\(1 \to 2 \to 4 \to 6\).\(6-4=2\)在前面出现过,\(4-2=2\)在前面出现过。那么根据这条additional chain计算6次幂的时候，可以从1次幂出发,用1次幂乘1次幂得到2次幂,再乘2次幂得到4次幂，再乘2次幂得到6次幂。

很可惜，对于数\(k\)求出得到\(k\)的最短additional chain是NP-hard的。但是有很好的近似算法。近似算法基于BFS。每次我们对于队头的数\(x\),枚举它对应的additional chain中的数\(y\),如果\(x+y\)还没有访问过那么将其入队,并将\(x\)对应的链后面接上\(x+y\). 这个预处理是\(O(k)\)的,且对快速幂的常数优化很显著。

如果\(k\)很大,比如\(10^{10000}\),可以采用十进制快速幂。但是用Method of Four Russians(俗称四毛子算法),可以将乘法次数减少到\(\log_2n+O(\frac{\log n}{\log \log n})\).具体方法见2017年国家集训队论文《非常规大小分块算法初探》

FFT的强常数优化