FFT（快速傅里叶变换）：HDU 5307 He is Flying

  1 #include <iostream>
  2 #include <cstring>
  3 #include <cstdio>
  4 #include <cmath>
  5 using namespace std;
  6 const long double PI=acos(-1.0);
  7 const int maxn=300010;
  8 struct complex{
  9     long double r,i;
 10     complex(long double r_=0.0,long double i_=0.0){
 11         r=r_;i=i_;
 12     }
 13     complex operator +(complex &a){
 14         return complex(r+a.r,i+a.i);
 15     }
 16     complex operator -(complex &a){
 17         return complex(r-a.r,i-a.i);
 18     }
 19     complex operator *(complex &a){
 20         return complex(r*a.r-i*a.i,r*a.i+i*a.r);
 21     }
 22 }A[maxn],B[maxn],C[maxn],D[maxn];
 23 
 24 void Rader(complex *a,int len){
 25     for(int i=1,j=len>>1;i<len-1;i++){
 26         if(i<j)swap(a[i],a[j]);
 27         int k=len>>1;
 28         while(j>=k){
 29             j-=k;
 30             k>>=1;
 31         }
 32         j+=k;
 33     }
 34 }
 35 
 36 void FFT(complex *a,int len,int on){
 37     Rader(a,len);
 38     for(int h=2;h<=len;h<<=1){
 39         complex wn(cos(-on*PI*2/h),sin(-on*PI*2/h));
 40         for(int j=0;j<len;j+=h){
 41             complex w(1.0,0);
 42             for(int k=j;k<j+(h>>1);k++){
 43                 complex x=a[k];
 44                 complex y=a[k+(h>>1)]*w;
 45                 a[k]=x+y;
 46                 a[k+(h>>1)]=x-y;
 47                 w=w*wn;
 48             }    
 49         }
 50     }
 51     if(on==-1)
 52         for(int i=0;i<len;i++)
 53             a[i].r/=len;
 54 }
 55 
 56 int d[maxn];
 57 int main(){
 58 #ifndef ONLINE_JUDGE
 59     //freopen();
 60     //freopen();
 61 #endif
 62     int T,n,len;
 63     long long tot,zero;
 64     scanf("%d",&T);
 65     while(T--){
 66         scanf("%d",&n);
 67         tot=zero=0;
 68         for(int i=1;i<=n;i++){
 69             scanf("%d",&d[i]);
 70             if(d[i]==0){
 71                 tot+=1;
 72                 zero+=tot*(tot+1)>>1;
 73             }
 74             else tot=0;
 75             d[i]+=d[i-1];
 76         }
 77         len=1;
 78         while(len<=2*d[n])len<<=1;
 79         memset(A,0,sizeof(A));
 80         memset(B,0,sizeof(B));
 81         memset(C,0,sizeof(C));
 82         memset(D,0,sizeof(D));
 83         for(int i=0;i<=n;i++){
 84             A[d[i]].r+=1.0;
 85             B[d[n]-d[i]].r+=i;
 86             C[d[i]].r+=i;
 87             D[d[n]-d[i]].r+=1.0;
 88         }
 89         FFT(A,len,1);FFT(B,len,1);
 90         for(int i=0;i<len;i++)
 91             A[i]=A[i]*B[i];
 92         FFT(A,len,-1);    
 93         FFT(C,len,1);FFT(D,len,1);
 94         for(int i=0;i<len;i++)
 95             C[i]=C[i]*D[i];
 96         FFT(C,len,-1);    
 97         printf("%lld\n",zero);
 98         for(int i=d[n]-1;i>=0;i--)
 99             printf("%lld\n",(long long)(A[i].r-C[i].r+0.5));        
100     }
101     return 0;
102 }