【CTSC2010】性能优化

mrsrz 12月 31, 2019

循环卷积，快速傅里叶变换

题目大意：

给定长度为 $n$ 的数列 $A,B$ 和正整数 $C$。

求 $A\times B^C$ 在循环卷积下的答案。

系数对质数 $n+1$ 取模，保证 $n$ 只包含 $2,3,5,7$ 这些质因数。

题解：

$\omega_n^n=\omega_n^0$，DFT 本身就是循环卷积。

所以我们对原数列进行长度为 $n$ 的 DFT，然后直接点值快速幂，再 IDFT 回来就是循环卷积后的结果。

考虑 FFT 算法本身的实质，它相当于每次把原序列分成均等的两部分，然后进行分治，再进行合并。所以才需要把序列长度扩充为 $2^n$。

$F(x)=a_0+a_1x+a_2x^2+\cdots+a_{n-1}x^{n-1}\\ F^{(0)}(x)=a_0+a_2^x+\cdots+a_{n-2}x^{\frac{n-2}2}\\ F^{(1)}(x)=a_1+a_3^x+\cdots+a_{n-1}x^{\frac{n-2}2}\\ F(x)=F^{(0)}(x^2)+xF^{(1)}(x^2)$

带入 $n$ 次单位根进行分治。

我们做长度为 $n$ 的 DFT，也需要每次分成均等的若干部分。

由于保证 $n$ 的质因子只有 $2,3,5,7$，所以每层最多分为 $7$ 份进行递归。

时间复杂度可以认为是 $O(n\log n)$。

Code：

#include<iostream>
#include<cstring>
using namespace std;
typedef long long LL;
const int N=524288;
int n,K,md,A[N],B[N],g,tmp[N];
inline int pow(int a,int b){
    int ret=1;
    for(;b;b>>=1,a=(LL)a*a%md)if(b&1)ret=(LL)ret*a%md;
    return ret;
}
int find(){
    for(int g=2;;++g){
        if(n%2==0&&pow(g,n/2)==1)continue;
        if(n%3==0&&pow(g,n/3)==1)continue;
        if(n%5==0&&pow(g,n/5)==1)continue;
        if(n%7==0&&pow(g,n/7)==1)continue;
        return g;
    }
}
void FFT(int*A,int n,const int g){
    if(n==1)return;
    const int bk=(n%2==0?2:(n%3==0?3:(n%5==0?5:7))),S=n/bk;
    int it=0;
    for(int K=0;K<bk;++K)
    for(int i=K;i<n;i+=bk)
    tmp[it++]=A[i];
    memcpy(A,tmp,sizeof(int)*n);
    const int g_=pow(g,bk);
    for(int i=0;i<n;i+=S)
    FFT(A+i,S,g_);
    memset(tmp,0,sizeof(int)*n);
    for(int i=0;i<n;++i){
        const int gi=pow(g,i);
        for(int k=0,G=1;k<bk;++k,G=(LL)G*gi%md)
        tmp[i]=(tmp[i]+(LL)G*A[i%S+S*k])%md;
    }
    memcpy(A,tmp,sizeof(int)*n);
}
int main(){
    ios::sync_with_stdio(0),cin.tie(0),cout.tie(0);
    cin>>n>>K,md=n+1;
    g=find();
    for(int i=0;i<n;++i)cin>>A[i];
    for(int i=0;i<n;++i)cin>>B[i];
    FFT(A,n,g),FFT(B,n,g);
    for(int i=0;i<n;++i)A[i]=(LL)A[i]*pow(B[i],K)%md;
    FFT(A,n,pow(g,md-2));
    const int iv=pow(n,md-2);
    for(int i=0;i<n;++i)A[i]=(LL)A[i]*iv%md;
    for(int i=0;i<n;++i)
    cout<<A[i]<<'\n';
    return 0;
}