时间序列的傅里叶变换实现 python 傅里叶变换时间复杂度

转载

AI大梦想家 2024-01-16 16:50:56

文章标签 时间序列的傅里叶变换实现 python i++ 多项式 #include 文章分类 Python 后端开发

　　快速傅里叶变换，是求两个多项式卷积的算法，其时间复杂度为$O(n\log n)$，优于普通卷积求法，且根据有关证明，快速傅里叶变换是基于变换求卷积的理论最快算法。

　　关于FFT的介绍，最详细易懂的是《算法导论》上的内容。

1.FFT&NTT模板

1 #include<cmath>
 2 #include<cstdio>
 3 #include<algorithm>
 4 #define rep(i,l,r) for (int i=l; i<=r; i++)
 5 using namespace std;
 6 
 7 const int N=300010;
 8 const double pi=acos(-1.);
 9 int n,m,L,x,rev[N];
10 
11 struct C{
12     double x,y;
13     C(double _x=0,double _y=0):x(_x),y(_y){}
14     C operator +(C &b){ return C(x+b.x,y+b.y); }
15     C operator -(C &b){ return C(x-b.x,y-b.y); }
16     C operator *(C &b){ return C(x*b.x-y*b.y,x*b.y+b.x*y); }
17 }a[N],b[N];
18 
19 void DFT(C a[],int n,bool f){
20     for (int i=0; i<n; i++) if (i<rev[i]) swap(a[i],a[rev[i]]);
21     for (int i=1; i<n; i<<=1){
22         C wn=C(cos(pi/i),(f?1:-1)*sin(pi/i));
23         for (int p=i<<1,j=0; j<n; j+=p){
24             C w(1,0);
25             for (int k=0; k<i; k++,w=w*wn){
26                 C x=a[j+k],y=w*a[i+j+k]; a[j+k]=x+y; a[i+j+k]=x-y;
27             }
28         }
29     }
30     if (!f) for (int i=0; i<n; i++) a[i].x/=n;
31 }
32 
33 int main(){
34     freopen("Dft.in","r",stdin);
35     freopen("Dft.out","w",stdout);
36     scanf("%d%d",&n,&m);
37     rep(i,0,n) scanf("%d",&x),a[i]=x;
38     rep(i,0,m) scanf("%d",&x),b[i]=x;
39     rep(i,0,max(n,m)) a[i]=C(a[i].x+b[i].x,a[i].x-b[i].x);
40     m=n+m; for (n=1; n<=m; n<<=1) L++;
41     rep(i,0,n-1) rev[i]=(rev[i>>1]>>1)|((i&1)<<(L-1));
42     DFT(a,n,1); rep(i,0,n-1) a[i]=a[i]*a[i]; DFT(a,n,0);
43     rep(i,0,m) printf("%d ",(int)(a[i].x/4+0.5));
44     return 0;
45 }

FFT

1 void NTT(int a[],int n,bool f){
 2     for (int i=0; i<n; i++) if (i<rev[i]) swap(a[i],a[rev[i]]);
 3     for (int i=1; i<n; i<<=1){
 4         int wn=ksm(3,f ? (mod-1)/(i<<1) : (mod-1)-(mod-1)/(i<<1));
 5         for (int p=i<<1,j=0; j<n; j+=p){
 6             int w=1;
 7             for (int k=0; k<i; k++,w=1ll*w*wn%mod){
 8                 int x=a[j+k],y=1ll*w*a[i+j+k]%mod;
 9                 a[j+k]=(x+y)%mod; a[i+j+k]=(x-y+mod)%mod;
10             }
11         }
12     }
13     if (f) return;
14     int inv=ksm(n,mod-2);
15     for (int i=0; i<n; i++) a[i]=1ll*a[i]*inv%mod;
16 }

NTT

1 #include<cmath>
  2 #include<cstdio>
  3 #include<algorithm>
  4 #include<cstring>
  5 #define mem(a) memset(a,0,sizeof(a))
  6 #define rep(i,l,r) for (int i=(l); i<=(r); i++)
  7 using namespace std;
  8 
  9 const int N=500010,mod=998244353,inv2=(mod+1)/2;
 10 int n,k,rev[N],inv[N],X[N],Y[N],A[N],B[N],C[N],D[N],E[N],F[N],G[N];
 11 
 12 void Print(int a[],int n=::n){ for (int i=0; i<n; i++) printf("%d ",a[i]); puts(""); }
 13 
 14 int ksm(int a,int b){
 15     int res=1;
 16     for (; b; a=1ll*a*a%mod,b>>=1)
 17         if (b & 1) res=1ll*res*a%mod;
 18     return res;
 19 }
 20 
 21 void NTT(int a[],int n,bool f){
 22     for (int i=0; i<n; i++) if (i<rev[i]) swap(a[i],a[rev[i]]);
 23     for (int i=1; i<n; i<<=1){
 24         int wn=ksm(3,f ? (mod-1)/(i<<1) : (mod-1)-(mod-1)/(i<<1));
 25         for (int p=i<<1,j=0; j<n; j+=p){
 26             int w=1;
 27             for (int k=0; k<i; k++,w=1ll*w*wn%mod){
 28                 int x=a[j+k],y=1ll*w*a[i+j+k]%mod;
 29                 a[j+k]=(x+y)%mod; a[i+j+k]=(x-y+mod)%mod;
 30             }
 31         }
 32     }
 33     if (f) return;
 34     int inv=ksm(n,mod-2);
 35     for (int i=0; i<n; i++) a[i]=1ll*a[i]*inv%mod;
 36 }
 37 
 38 void mul(int a[],int b[],int l){
 39     int n=1,L=0;
 40     for (; n<(l<<1); n<<=1) L++;
 41     for (int i=0; i<n; i++) rev[i]=(rev[i>>1]>>1)|((i&1)<<(L-1));
 42     NTT(a,n,1); NTT(b,n,1);
 43     for (int i=0; i<n; i++) a[i]=1ll*a[i]*b[i]%mod;
 44     NTT(a,n,0); NTT(b,n,0);
 45 }
 46 
 47 void Inv(int a[],int b[],int l){
 48     if (l==1){ b[0]=ksm(a[0],mod-2); return; }
 49     Inv(a,b,l>>1); int n=1,L=0;
 50     for (; n<(l<<1); n<<=1) L++;
 51     for (int i=0; i<n; i++) rev[i]=(rev[i>>1]>>1)|((i&1)<<(L-1));
 52     for (int i=0; i<l; i++) A[i]=a[i];
 53     NTT(A,n,1); NTT(b,n,1);
 54     for (int i=0; i<n; i++) b[i]=1ll*b[i]*(2-1ll*A[i]*b[i]%mod+mod)%mod;
 55     NTT(b,n,0);
 56     for (int i=l; i<n; i++) b[i]=0;
 57     for (int i=0; i<n; i++) A[i]=0;
 58 }
 59 
 60 void Sqrt(int a[],int b[],int l){
 61     if (l==1){ b[0]=sqrt(a[0]); return; }
 62     Sqrt(a,b,l>>1); Inv(b,B,l); int n=1,L=0;
 63     for (; n<(l<<1); n<<=1) L++;
 64     for (int i=0; i<n; i++) rev[i]=(rev[i>>1]>>1)|((i&1)<<(L-1));
 65     for (int i=0; i<l; i++) C[i]=a[i];
 66     NTT(C,n,1); NTT(B,n,1); NTT(b,n,1);
 67     for (int i=0; i<n; i++) b[i]=1ll*inv2*(b[i]+1ll*C[i]*B[i]%mod)%mod;
 68     NTT(b,n,0);
 69     for (int i=l; i<n; i++) b[i]=0;
 70     for (int i=0; i<n; i++) C[i]=B[i]=0;
 71 }
 72 
 73 void Deri(int a[],int b[],int l){
 74     for (int i=1; i<l; i++) b[i-1]=1ll*i*a[i]%mod;
 75 }
 76 
 77 void Inte(int a[],int b[],int l){
 78     for (int i=1; i<l; i++) b[i]=1ll*a[i-1]*inv[i]%mod; b[0]=0;
 79 }
 80 
 81 void Ln(int a[],int b[],int l){
 82     Deri(a,D,l); Inv(a,E,l); mul(D,E,l); Inte(D,b,l);
 83     for (int i=0; i<(l<<1); i++) D[i]=E[i]=0;
 84 }
 85 
 86 void Exp(int a[],int b[],int l){
 87     if (l==1){ b[0]=1; return; }
 88     Exp(a,b,l>>1); Ln(b,F,l); int n=1,L=0;
 89     for (; n<(l<<1); n<<=1) L++;
 90     for (int i=0; i<n; i++) rev[i]=(rev[i>>1]>>1)|((i&1)<<(L-1));
 91     for (int i=0; i<l; i++) F[i]=(-F[i]+a[i]+mod)%mod; F[0]=(F[0]+1)%mod;
 92     NTT(F,n,1); NTT(b,n,1);
 93     for (int i=0; i<n; i++) b[i]=1ll*b[i]*F[i]%mod;
 94     NTT(b,n,0);
 95     for (int i=l; i<n; i++) b[i]=0;
 96     for (int i=0; i<n; i++) F[i]=0;
 97 }
 98 
 99 void Ksm(int a[],int b[],int k,int l){
100     Ln(a,G,l);
101     for (int i=0; i<l; i++) G[i]=1ll*k*G[i]%mod;
102     Exp(G,b,l);
103 }
104 
105 int main(){
106     freopen("polynomial.in","r",stdin);
107     freopen("polynomial.out","w",stdout);
108     scanf("%d%d",&n,&k);
109     for (int i=0; i<n; i++) scanf("%d",&X[i]);
110     int l=1; for (; l<=n; l<<=1); inv[0]=inv[1]=1;
111     rep(i,2,l) inv[i]=1ll*(mod-mod/i)*inv[mod%i]%mod;
112     Sqrt(X,Y,l); mem(X); //Print(Y);
113     Inv(Y,X,l); mem(Y); //Print(X);
114     Inte(X,Y,l); mem(X); //Print(Y);
115     Exp(Y,X,l); mem(Y); //Print(X);
116     Inv(X,Y,l); Y[0]=(Y[0]+1)%mod; mem(X); //Print(Y);
117     Ln(Y,X,l); X[0]=(X[0]+1)%mod; mem(Y); //Print(X);
118     Ksm(X,Y,k,l); mem(X); //Print(Y);
119     Deri(Y,X,n); mem(Y); Print(X);
120     return 0;
121 }

多项式全家桶(COGS2189)

2.牛顿迭代法

现已知$G(F(x))\equiv 0\ (\text{mod}\ x^n)$，考虑如何倍增求出$F(x)$。
我们先将x扩展到2的次幂项（多余位补0），方便后面的倍增。
当n=1时，直接算出常数项即可。
否则，先递归下去求出$F_0(x)$使得$G(F_0(x))\equiv 0\ (\text{mod}\ x^\frac{n}{2})$。
接着将$G(F(x))$在$F_0(x)$处泰勒展开，有：$$G(F(x))=G(F_0(x))+G'(F_0(x))(F(x)-F_0(x))+\frac{G''(F_0(x))}{2}(F(x)-F_0(x))^2+...$$由于$F(x)$与$F_0(x)$前$\frac{n}{2}$项都相同，所以上式从第三项开始也都是0。
于是有：$G(F(x))\equiv G(F_0(x))+G'(F_0(x))(F(x)-F_0(x))\equiv 0(\text{mod}\ x^n)$
移项得：$F(x)=F_0(x)-\frac{G(F_0(x))}{G'(F_0(x))}\ (\text{mod}\ x^n)$
（顺便一提，这里的$F_0(x)$只是一个多项式，可以当作普通泰勒展开中的$x$看待，不要混淆概念将$G'(F_0(x))$链式法则了）
这样，我们递归下去求出$F_0(x)$，再多项式求逆与乘法，就得到了$F(x)$。
时间复杂度：$T(n)=T(n/2)+O(n\log n)=O(n\log n)$

3.多项式求逆

利用上面牛顿迭代法，考虑如何求出一个多项式在模意义下的逆元，即对于多项式$A(x)$，求出$B(x)$使得$A(x)B(x)\equiv\ 1(\text{mod}\ x^n)$（当然这里不用牛顿迭代法也能得出相同的结论）。
构造函数$G(B(x))=A(x)B(x)-1$，则有$G(B(x))\equiv 0\ (\text{mod}\ x^n)$
同样将次数界扩展到2的次幂，然后先递归下去求出$B_0(x)$使得$A(x)B_0(x)\equiv 1\ (\text{mod}\ x^\frac{n}{2})$。
接下来暴力带入牛顿迭代的式子：$B(x)=B_0(x)-\frac{G(B_0(x))}{G'(B_0(x))}=B_0(x)-\frac{A(x)B_0(x)-1}{A(x)}=B_0(x)-B_0(x)(A(x)B_0(x)-1)=B_0(x)(2-A(x)B_0(x))$

1 void Inv(int a[],int b[],int l){
 2     if (l==1){ b[0]=ksm(a[0],mod-2); return; }
 3     Inv(a,b,l>>1); int n=1,L=0;
 4     for (; n<(l<<1); n<<=1) L++;
 5     for (int i=0; i<n; i++) rev[i]=(rev[i>>1]>>1)|((i&1)<<(L-1));
 6     for (int i=0; i<l; i++) A[i]=a[i];
 7     NTT(A,n,1); NTT(b,n,1);
 8     for (int i=0; i<n; i++) b[i]=1ll*b[i]*(2-1ll*A[i]*b[i]%mod+mod)%mod;
 9     NTT(b,n,0);
10     for (int i=l; i<n; i++) b[i]=0;
11     for (int i=0; i<n; i++) A[i]=0;
12 }

Inv

4.多项式开根

和求逆一样套牛顿迭代的式子：设$G(B(x))=B^2(x)-A(x)$，则递归下去求出$B_0(x)$，再$B(x)\equiv B_0(x)-\frac{G(B_0(x))}{G'(B_0(x))}=B_0(x)-\frac{B_0^2(x)-A(x)}{2B_0(x)}=\frac{B_0^2(x)+A(x)}{2B_0(x)}\ (\text{mod}\ x^n)$

1 void Sqrt(int a[],int b[],int l){
 2     if (l==1){ b[0]=sqrt(a[0]); return; }
 3     Sqrt(a,b,l>>1); Inv(b,B,l); int n=1,L=0;
 4     for (; n<(l<<1); n<<=1) L++;
 5     for (int i=0; i<n; i++) rev[i]=(rev[i>>1]>>1)|((i&1)<<(L-1));
 6     for (int i=0; i<l; i++) C[i]=a[i];
 7     NTT(C,n,1); NTT(B,n,1); NTT(b,n,1);
 8     for (int i=0; i<n; i++) b[i]=1ll*inv2*(b[i]+1ll*C[i]*B[i]%mod)%mod;
 9     NTT(b,n,0);
10     for (int i=l; i<n; i++) b[i]=0;
11     for (int i=0; i<n; i++) C[i]=B[i]=0;
12 }

Sqrt

5.多项式求导与积分

这个很简单，直接按定义来即可，复杂度线性。

1 void Deri(int a[],int b[],int l){
2     for (int i=1; i<l; i++) b[i-1]=1ll*i*a[i]%mod;
3 }
4 
5 void Inte(int a[],int b[],int l){
6     for (int i=1; i<l; i++) b[i]=1ll*a[i-1]*inv[i]%mod; b[0]=0;
7 }

Derivation&Integration

6.多项式求ln

这个不需要牛顿迭代，因为设$\ln(A(x))=B(x)$，求导得$\int\frac{A'(x)}{A(x)}=B(x)$，于是做一遍多项式求导、一遍求逆、一遍乘法和一遍积分即可。要求常数项为1。

1 void Ln(int a[],int b[],int l){
2     Deri(a,D,l); Inv(a,E,l); mul(D,E,l); Inte(D,b,l);
3     for (int i=0; i<(l<<1); i++) D[i]=E[i]=0;
4 }

7.多项式求exp

$exp(A(x))=B(x)$，两边求导得$A(x)=\ln(B(x))$，设$G(B(x))=\ln(B(x))-A(x)$，递归下去求出$B_0(x)$，再$B(x)=B_0(x)-\frac{G(B_0(x))}{G'(B_0(x))}=B_0(x)-\frac{\ln(B(x))-A(x)}{\frac{1}{B_0(x)}}$=$B_0(x)(1-\ln(B(x))+A(x))$。要求常数项为0。

1 void Exp(int a[],int b[],int l){
 2     if (l==1){ b[0]=1; return; }
 3     Exp(a,b,l>>1); Ln(b,F,l); int n=1,L=0;
 4     for (; n<(l<<1); n<<=1) L++;
 5     for (int i=0; i<n; i++) rev[i]=(rev[i>>1]>>1)|((i&1)<<(L-1));
 6     for (int i=0; i<l; i++) F[i]=(-F[i]+a[i]+mod)%mod; F[0]=(F[0]+1)%mod;
 7     NTT(F,n,1); NTT(b,n,1);
 8     for (int i=0; i<n; i++) b[i]=1ll*b[i]*F[i]%mod;
 9     NTT(b,n,0);
10     for (int i=l; i<n; i++) b[i]=0;
11     for (int i=0; i<n; i++) F[i]=0;
12 }

Exp

8.多项式快速幂

$A^k(x)=exp(k*\ln(A(x)))$

不管怎么套都是一个log，但显然常数巨大。

1 void Ksm(int a[],int b[],int k,int l){
2     Ln(a,G,l);
3     for (int i=0; i<l; i++) G[i]=1ll*k*G[i]%mod;
4     Exp(G,b,l);
5 }

Pow

9.多项式除法

要求的就是给定两个多项式$A(x)$,$B(x)$，其项数为$n$,$m$
求解一个$n-m$项的多项式$C(x)$，以及一个小于$n-m$项的多项式$R(x)$。
满足：$A(x)=B(x)*C(x)+R(x)$。
定义一个操作$R$,其中$R$就是$Reverse$，$A^R(x)=x^nA(\frac{1}{x})$。这个操作说白点就是$A(x)[x^i]$对应$A^R(x)[x^{n-i}]$，也就是把所有的系数给翻转过来。
然后推式子：
$$\begin{aligned} A(x)&=B(x)*C(x)+R(x)\\ A(\frac{1}{x})&=B(\frac{1}{x})*C(\frac{1}{x})+R(\frac{1}{x})\\ x^nA(\frac{1}{x})&=(x^m*B(\frac{1}{x}))*(x^{n-m}*C(\frac{1}{x}))+x^nR(\frac{1}{x})\\ A^R(x)&=B^R(x)*C^R(x)+R^R(x)*x^{n-m+1} \end{aligned}$$
到了这里我们把等号换成同余，把整个式子在模$x^{n-m+1}$意义下进行。
$$\begin{aligned} A^R(x)&\equiv B^R(x)*C^R(x)+R^R(x)*x^{n-m+1}\\ &\equiv B^R(x)*C^R(x) \end{aligned}$$
所以在模意义下，我们可以利用多项式求逆求解$C^R(x)=\frac{A^R(x)}{B^R(x)}mod\ x^{n-m+1}$
那么就有$R(x)=A(x)-B(x)*C(x)$。

10.分治FFT

分治FFT是求这样一类卷积：现已知$g[0],g[1],...,g[n-1]$和$f[0]$，$f[i]=\sum\limits_{j=1}^{i}f[i-j]g[j]$，求$f[1],...,f[n]$。
由于式子左右都有$f$，需要用类似$CDQ$分治的做法来实现。复杂度$O(n\log^2 n)$
假设我们现在需要求$f[l],...,f[r]$先递归到左边处理出$f[l],...,f[mid]$，再计算左半边对右半边的影响。这里直接卷上$g$即可，主要要作适当移位。（当然如果g直接给出也是可以直接多项式求逆的）

1 void CDQ(int l,int r){
 2     if (l==r) return;
 3     int mid=(l+r)>>1,lim=r-l+1,n=1,L=0;
 4     CDQ(l,mid);
 5     while (n<lim) n<<=1,L++;
 6     for (int i=0; i<n; i++) rev[i]=(rev[i>>1]>>1)|((i&1)<<(L-1));
 7     for (int i=0; i<n; i++) a[i]=b[i]=0;
 8     rep(i,l,mid) a[i-l]=f[i];
 9     rep(i,0,r-l) b[i]=g[i];
10     NTT(a,n,1); NTT(b,n,1);
11     for (int i=0; i<n; i++) a[i]=1ll*a[i]*b[i]%mod;
12     NTT(a,n,0);
13     rep(i,mid+1,r) f[i]=(f[i]+a[i-l])%mod;
14     CDQ(mid+1,r);
15 }

CDQ

11.MTT

任意模数FFT，在模数不为998244353，直接做FFT又可能炸精度的情况下使用。

有两种写法，一种是myy的三模数FFT，一种是下面的做7次FFT的做法。

求$F(x)$与$G(x)$在任意模数下的卷积。
为什么不能直接FFT乘然后再取模？因为直接乘结果会爆long long。
考虑拆系数。设一个常数M，把$F(x)$和$G(x)$拆成：$A(x)=\frac{F(x)}M$，$B(x)=F(x)\%M$，$C(x)=\frac{G(x)}M$，$D(x)=G(x)\%M$。
这样答案就变成了：
$M^2A(x)C(x)+M(A(x)D(x)+B(x)C(x))+B(x)D(x)$
直接FFT就不会爆long long了。$M$取$\sqrt {mod}$（我取了32768）最优。
这种方法一共要做7次DFT。

1 void calc(int a1[],int a2[],int l){
 2     int n=1,L=0;
 3     for (; n<l; n<<=1,L++);
 4     for (int i=0; i<n; i++) rev[i]=(rev[i>>1]>>1)|((i&1)<<(L-1));
 5     for (int i=0; i<n; i++) A[i]=a1[i]>>15,B[i]=a1[i]&0x7fff;
 6     for (int i=0; i<n; i++) C[i]=a2[i]>>15,D[i]=a2[i]&0x7fff;
 7     FFT(A,n,1); FFT(B,n,1); FFT(C,n,1); FFT(D,n,1);
 8     for (int i=0; i<n; i++) s1[i]=A[i]*C[i],s2[i]=A[i]*D[i]+B[i]*C[i],s3[i]=B[i]*D[i];
 9     FFT(s1,n,-1); FFT(s2,n,-1); FFT(s3,n,-1);
10     for (int i=0; i<n; i++)
11         a1[i]=((((ll)round(s1[i].x)%mod)<<30)%mod+(((ll)round(s2[i].x)%mod)<<15)%mod+(ll)round(s3[i].x)%mod)%mod;
12 }

MTT

12.应用

1.加速卷积运算

卷积运算一般会用在容斥、DP等题目中，所有形如$c_{k}=\sum_{i=0}^{k}a_{i}b_{k-i}$的计算都可以用FFT加速。同时，形如$c_{k}=\sum_{i=k}^{n-1}a_{i}b_{i-k}$的问题也可以，注意到只需要将$b[]$翻转一下做一次卷积，然后$a[n],...,a[2n-1]$的位置存的就是$c[0],...,c[n-1]$的答案了，同时分治FFT也能解决相同形式的问题。

另外，在快速求解第一、第二类斯特林数时也有应用

例题：HDU4609,BZOJ5306,BZOJ4555

2.加速字符串匹配

比如要求对一个字符串的每一个位置，统计以它为轴的最长回文子序列长度。

对字符集里的每种字符做一遍，所有等于当前处理字符的位置赋为1，否则赋为0。以x轴回文意味着s[x-i]=s[x+i]，注意到这时s[x-i]*s[x+i]为1，且如果做卷积，这个值会被统计到一个固定位置x-i+x+i=2x处。所以做字符集大小次FFT即可。

同理，普通字符串匹配也可以类似地做，但卷积后每个位置记录的就是固定的差的信息了。

关于带通配符的字符串匹配问题也有FFT应用：BZOJ3160,BZOJ4503。

3.优化DP转移

一个比较常见的套路是，需要求一个满足题设条件的DP数组f[]，可以先放宽题设限制求一个数组g[]，然后讲f和g都视为生成函数，根据二者的关系做FFT。这种方法往往出现于带有组合数、容斥的DP中。例题：BZOJ3456,Luogu5162

4.分治+FFT

求解形如$\prod(x^{k_i}+a_i)$且$\sum k_i$在1e5左右范围内时可以应用的一种做法。

分治，递归到两边分别求解，然后一次卷积合并。$O(n\log^2 n)$