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

A和B的運算是卷積形式

考慮把A和B DFT

D(A)∗D(B)=0,D(B)≠0$D(A)\ast D(B)=0,D(B)\ne 0$

也就是 A DFT後至少有一個0

求出模某個質數意義下的 2n$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;
}