HDU 5307 He is Flying (生成函数+FFT)

题目传送门

题目大意：给你一个长度为$n$的自然数序列$a$，定义一段区间的权值为这一段区间里所有数的和，分别输出权值为$[0,\sum a_{i}]$的区间的长度之和

想到了生成函数的话，这道题并不难做。但很多细节真是不太好搞

我们首先预处理出前缀和s，那么一段区间$[l,r]$的权值就是$s_{r}-s_{l-1}$

容易联想到卷积

第一个多项式是 区间右端点的前缀和 作为指数的生成函数，每一项的系数是 右端点的编号之和

第二个多项式是 区间左端点的前缀和 作为指数的生成函数，每一项的系数是 左端点的编号之和

然而区间长度是相减而不是相乘

我们可以把问题转化成 右端点编号$*1$-左端点编号$*1$，求两次卷积再相减即可

然而左端点的前缀和是负数，我们把生成函数整体右移

然而序列里还有$0$的情况

如果序列里出现了连续的$0$，我们发现这部分答案我们无法通过卷积统计

因为按照我们的方法，在多项式对应的相同的位置卷积的话，两次统计的答案就被减掉了

所以连续的$0$对答案的影响通过$O(n)$扫一遍统计

每新加入一个新的$0$，就会多产生一个等差数列的贡献

另外答案比较大，$FFT$需要开$long\;double$

 #include <cmath>
 #include <cstdio>
 #include <cstring>
 #include <algorithm>
 #define N1 (1<<18)
 #define M1 (N1<<1)
 #define il inline
 #define dd double
 #define ld long double
 #define ll long long
 using namespace std;

 int T,n;

 int gint()
 {
     int ret=,fh=;char c=getchar();
     while(c<''||c>''){if(c=='-')fh=-;c=getchar();}
     while(c>=''&&c<=''){ret=ret*+c-'';c=getchar();}
     return ret*fh;
 }

 const ld pi=acos(-);
 struct cp{
 ld x,y;
 friend cp operator + (const cp &s1,const cp &s2){ return (cp){s1.x+s2.x,s1.y+s2.y}; }
 friend cp operator - (const cp &s1,const cp &s2){ return (cp){s1.x-s2.x,s1.y-s2.y}; }
 friend cp operator * (const cp &s1,const cp &s2){ return (cp){s1.x*s2.x-s1.y*s2.y,s1.y*s2.x+s1.x*s2.y}; }
 }a[N1],b[N1],c[N1];
 int r[N1];
 void FFT(cp *s,int len,int type)
 {
     int i,j,k; cp wn,w,t;
     for(i=;i<len;i++) if(i<r[i]) swap(s[i],s[r[i]]);
     for(k=;k<=len;k<<=)
     {
         wn=(cp){cos(pi*2.0*type/k),sin(pi*2.0*type/k)};
         for(i=;i<len;i+=k)
         {
             w=(cp){,};
             for(j=;j<(k>>);j++,w=w*wn)
             {
                 t=w*s[i+j+(k>>)];
                 s[i+j+(k>>)]=s[i+j]-t;
                 s[i+j]=s[i+j]+t;
             }
         }
     }
 }
 void FFT_Main(int len)
 {
     FFT(a,len,); FFT(b,len,);
     for(int i=;i<len;i++) c[i]=a[i]*b[i];
     FFT(c,len,-);
     for(int i=;i<len;i++) c[i].x/=len;
 }

 int v[N1],s[N1];
 ll ans[N1];

 int main()
 {
     scanf("%d",&T);
     while(T--) {

     memset(v,,sizeof(v)); memset(s,,sizeof(s)); memset(r,,sizeof(r)); memset(ans,,sizeof(ans));
     int i,j,maxn=,len,L,num;
     scanf("%d",&n);
     for(i=;i<=n;i++) v[i]=gint(), s[i]=s[i-]+v[i], maxn=max(maxn,s[i]);
     if(!maxn)
     {
         for(i=;i<=n;i++)
             ans[]+=1ll*(i+)*i/;
         printf("%lld\n",ans[]);
         continue;
     }

     for(len=,L=;len<maxn+maxn+;len<<=,L++);
     for(i=;i<len;i++) r[i]=(r[i>>]>>)|((i&)<<(L-));

     memset(a,,sizeof(a)); memset(b,,sizeof(b));
     for(i=;i<=n;i++) a[s[i]].x+=i;
     for(i=;i<=n;i++) b[-s[i]+maxn].x++;
     FFT_Main(len);
     for(i=maxn;i<=(maxn<<);i++) ans[i-maxn]=(ll)(c[i].x+0.5);

     memset(a,,sizeof(a)); memset(b,,sizeof(b));
     for(i=;i<=n;i++) a[s[i]].x++;
     for(i=;i<=n;i++) b[-s[i]+maxn].x+=i;
     FFT_Main(len);
     for(i=maxn;i<=(maxn<<);i++) ans[i-maxn]-=(ll)(c[i].x+0.5);

     for(i=,num=;i<=n;i++)
         if(!v[i]){ num++; ans[]+=1ll*(num+)*(num)/2ll; }
         else{ num=; }
     for(i=;i<=maxn;i++) printf("%lld\n",ans[i]);
     //puts("");
     }
     return ;

 }