hdu 5307 He is Flying(推公式+FFT)

#include<cstdio>
#include<cstring>
#include<cmath>
#include<algorithm>
using namespace std;
#define ld long double
ld pi=acos(-1.0);   //注意π的精度
int num[500010],sum[500010];
long long snum[500010],ans[500010];
long long p,res;
struct Complex{
   ld r;
   ld i;
   Complex(){};
   Complex(ld a,ld b){
       r=a;i=b;
   }
   Complex operator +(const Complex &t)const{
      return Complex(r+t.r,i+t.i);
   }
   Complex operator -(const Complex &t)const{
      return Complex(r-t.r,i-t.i);
   }
   Complex operator *(const Complex &t)const{
      return Complex(r*t.r-i*t.i,r*t.i+i*t.r);
   }
}a1[500010],a2[500010];
void fft(Complex y[],int n,int rev){
    for(int i=1,j,k,t;i<n;i++){
        for(j=0,k=n>>1,t=i;k;k>>=1,t>>=1) j=j<<1|t&1;
        if(i<j) swap(y[i],y[j]);
    }
    for(int s=2,ds=1;s<=n;ds=s,s<<=1){
      Complex wn=Complex(cos(rev*2*pi/s),sin(rev*2*pi/s)),w=Complex(1,0),t;
      for(int k=0;k<ds;k++,w=w*wn){
        for(int i=k;i<n;i+=s){
            t=w*y[i+ds];
            y[i+ds]=y[i]-t;
            y[i]=y[i]+t;
        }
      }
    }
    if(rev==-1) for(int i=0;i<n;i++) y[i].r/=n;
}
int main()
{
    int i,j,k,t,n;
    snum[0]=0;
    for(long long i=1;i<=100000;i++)
        snum[i]=snum[i-1]+i*(i+1)/2;  //区间长度
    scanf("%d",&t);
    while(t--){
        memset(sum,0,sizeof(sum));
        scanf("%d",&n);
        for(i=0;i<n;i++){
            scanf("%d",&num[i]);
        }
        sum[0]=num[0];
        for(i=1;i<n;i++){
           sum[i]=sum[i-1]+num[i];  //区间和
        }
        p=0;res=0;
        for(i=0;i<n;i++){
          if(num[i]==0){
             p++;
          }
          else{
            res+=snum[p];
             p=0;
          }
        }
        res+=snum[p];
        printf("%lld\n",res);  //预处理区间和为0的情况
        int total=sum[n-1];    //多项式长度
        int total2=total*2;
        memset(a1,0,sizeof(a1));
        memset(a2,0,sizeof(a2));
        int len=1;
        while(len<=total2) len<<=1;
        for(i=0;i<n;i++){       //构造多项式
            a1[sum[i]].r+=i+1;
            if(i!=n-1)
                a2[total-sum[i]].r+=1;
        }
        a2[total].r+=1;
        fft(a1,len,1);
        fft(a2,len,1);
        for(i=0;i<=len;i++) a1[i]=a1[i]*a2[i];
        fft(a1,len,-1);
        for(i=0;i<=len;i++) ans[i]=(long long)(a1[i].r+0.5);

        memset(a1,0,sizeof(a1));
        memset(a2,0,sizeof(a2));
        for(i=0;i<n;i++){    //构造多项式
           a1[sum[i]].r+=1;
           if(i!=n-1){
             a2[total-sum[i]].r+=i+1;
           }
        }
        fft(a1,len,1);
        fft(a2,len,1);
        for(i=0;i<=len;i++) a1[i]=a1[i]*a2[i];
        fft(a1,len,-1);
        for(i=0;i<=len;i++) ans[i]-=(long long)(a1[i].r+0.5);
        for(i=total+1;i<=total2;i++)  //结果
            printf("%lld\n",ans[i]);
    }
    return 0;
}