快速傅里叶变换(FFT)、快速数论变换(NTT)模板

FFT

struct Z
{
    double x,y;
    Z(double _x=0,double _y=0):x(_x),y(_y){};
};
Z operator +(Z x,Z y) {return Z(x.x+y.x,x.y+y.y);}
Z operator -(Z x,Z y) {return Z(x.x-y.x,x.y-y.y);}
Z operator *(Z x,Z y) {return Z(x.x*y.x-x.y*y.y,x.x*y.y+x.y*y.x);}
int bit[M],n,m;
Z wi[M+1],a[M],b[M];
void prp()
{
    fo(i,0,M-1) bit[i]=(bit[i>>1]>>1)|((i&1)<<(L-1));
    fo(i,0,M) wi[i]=Z(cos(2*pi*i/M),sin(2*pi*i/M));
}
void DFT(Z *a,int pd)
{
    fo(i,0,M-1) if(i<bit[i]) swap(a[i],a[bit[i]]);
    int half=1,lim=(M>>1);
    Z v;
    for(int m=2;m<=M;half=m,m<<=1,lim>>=1)
    {
        fo(i,0,half-1)
        {
            Z w=(pd==1)?wi[i*lim]:wi[M-i*lim];
            for(int j=i;j<M;j+=m)
            {
                v=w*a[j+half];
                a[j+half]=a[j]-v;
                a[j]=a[j]+v;
            }
        }
    }
    if(pd<0) fo(i,0,M-1) a[i].x/=M;
}
int main()
{
    cin>>n>>m;
    fo(i,0,n) scanf("%lf",&a[i].x);
    fo(i,0,m) scanf("%lf",&b[i].x);
    prp();
    DFT(a,1),DFT(b,1);
    fo(i,0,M-1) a[i]=a[i]*b[i];
    DFT(a,-1);
    fo(i,0,n+m) 
    {
        int v=(int)(a[i].x+0.5);
        printf("%d ",v);
    }
}

NTT

inline LL md(LL x)
{
    return(x<0)?(x+mo):((x>=mo)?x-mo:x);
}
void NTT(LL *a,int pd,int num)
{
    prp(num);
    fo(i,0,num-1) if(i<bit[i]) swap(a[i],a[bit[i]]);
    int lim=num>>1,half=1;
    LL v;
    for(int m=2;m<=num;half=m,m<<=1,lim>>=1)
    {
        fo(i,0,half-1)
        {
            LL w=(pd==1)?wi[i*lim]:wi[num-i*lim];
            for(int j=i;j<num;j+=m)
            {
                v=(w*a[j+half])%mo;
                a[j+half]=md(a[j]-v);
                a[j]=md(a[j]+v);
            }
        }
    }
    if(pd<0) fo(i,0,num-1) a[i]=a[i]*ny[num]%mo;
}

猜你喜欢

转载自blog.csdn.net/hzj1054689699/article/details/80050733