Advertisement

[Bluestein's Algorithm DFT] Codechef REALSET. Petya and Sequence

阅读量:

A和B的运算是卷积形式

考虑把A和B DFT

D(A)*D(B)=0 , D(B)\neq0

也就是 A DFT后至少有一个0

求出模某个质数意义下的 2n 次单位根

用Bluestein’s Algorithm DFT

复制代码
    #include <cstdio>
    #include <iostream>
    #include <algorithm>
    #include <cmath>
    #include <ctime>
    #include <assert.h>
    
    using namespace std;
    
    const int N=300010;
    
    int n,r,pp[N];
    
    inline int Prime(int p){
      for(int i=2;i*i<=p;i++)
    if(p%i==0) return 0;
      return 1;
    }
    
    inline void Getdivisor(int p){
      pp[0]=0;
      for(int i=2;i*i<=p;i++)
    if(p%i==0){
      pp[++*pp]=i;
      while(p%i==0) p/=i;
    }
      if(p^1) pp[++*pp]=p;
    }
    
    inline int Pow(int x,int y,int p){
      int ret=1; x%=p;
      for(;y;y>>=1,x=1LL*x*x%p) if(y&1) ret=1LL*x*ret%p;
      return ret;
    }
    
    inline bool check(int x,int p){
      for(int i=1;i<=*pp;i++)
    if(Pow(x,(p-1)/pp[i],p)==1) return false;
      return true;
    }
    
    inline int Getroot(int p){
      for(int i=2;;i++)
    if(check(i,p)) return i;
      assert(0);
    }
    
    int rev[N];
    
    struct FFTer{
      int P,g;
      int w[2][N];
    
      inline void Pre(int n){
    int G=Pow(g,(P-1)/n,P);
    w[0][0]=w[1][0]=1;
    for(int i=1;i<n;i++) w[1][i]=1LL*w[1][i-1]*G%P;
    for(int i=1;i<n;i++) w[0][i]=w[1][n-i];
      }
    
      inline void FFT(int *a,int n,int r){
    for(int i=1;i<n;i++) if(rev[i]>i) swap(a[i],a[rev[i]]);
    for(int i=1;i<n;i<<=1)
      for(int j=0;j<n;j+=(i<<1))
    for(int k=0;k<i;k++){
      int x=a[j+k],y=1LL*a[j+k+i]*w[r][n/(i<<1)*k]%P;
      a[j+k]=(x+y)%P; a[j+k+i]=(x+P-y)%P;
    }
    if(!r) for(int i=0,inv=Pow(n,P-2,P);i<n;i++) a[i]=1LL*a[i]*inv%P;
      }
    }FFT[2];
    
    int A[N],B[N],C[N];
    int c[2][N],d[2][N];
    
    inline void Bluestein(int *a,int n,int p,int r){
      int m,L=0,omega=Pow(r,(p-1)/(2*n),p);
      for(m=1;m<n;m<<=1,L++); m<<=1;
      FFT[0].Pre(m); FFT[1].Pre(m);
      for(int i=1;i<m;i++) rev[i]=(rev[i>>1]>>1)|((i&1)<<L);
      for(int i=0;i<n;i++){
    A[i]=1LL*(a[i]+p)%p*Pow(omega,1LL*i*i%(2*n),p)%p;
    B[i]=Pow(omega,2*n-1LL*i*i%(2*n),p);
    if(i) B[m-i]=B[i];
      }
      for(int k=0;k<2;k++){
    int *cc=c[k],*dd=d[k],cp=FFT[k].P;
    for(int i=0;i<m;i++) 
      cc[i]=A[i],dd[i]=B[i];
    FFT[k].FFT(cc,m,1);
    //FFT[k].FFT(cc,m,0);
    FFT[k].FFT(dd,m,1);
    for(int i=0;i<m;i++)
      cc[i]=1LL*cc[i]*dd[i]%cp;
    FFT[k].FFT(cc,m,0);
      }
      int p0=FFT[0].P,p1=FFT[1].P;
      int inv=Pow(p0,p1-2,p1);
      for(int i=0;i<m;i++){
    int a0=c[0][i],a1=c[1][i];
    int k=1LL*(a1-a0%p1+p1)*inv%p1;
    C[i]=(1LL*k*p0+a0)%p;
    A[i]=B[i]=0;
      }
    }
    
    int a[N];
    
    int main(){
      freopen("1.in","r",stdin);
      freopen("1.out","w",stdout);
      int t; scanf("%d",&t);
      FFT[0].P=998244353; FFT[0].g=3;
      FFT[1].P=1005060097; FFT[1].g=5;
      while(t--){
    scanf("%d",&n);
    for(int i=0;i<n;i++) scanf("%d",&a[i]);
    int p=10000000/(2*n)*(2*n)+1;
    while(!Prime(p)) p+=2*n;
    Getdivisor(p-1);
    r=Getroot(p);
    Bluestein(a,n,p,r);
    int flg=0;
    for(int i=0;i<n;i++)
      if(C[i]==0) { flg=1; break; }
    puts(flg?"YES":"NO");
      }
      //printf("%d\n",clock());
      return 0;
    }

全部评论 (0)

还没有任何评论哟~