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=0,fh=1;char c=getchar();
    while(c<'0'||c>'9'){if(c=='-')fh=-1;c=getchar();}
    while(c>='0'&&c<='9'){ret=ret*10+c-'0';c=getchar();}
    return ret*fh;
}

const ld pi=acos(-1);
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=0;i<len;i++) if(i<r[i]) swap(s[i],s[r[i]]);
    for(k=2;k<=len;k<<=1)
    {
        wn=(cp){cos(pi*2.0*type/k),sin(pi*2.0*type/k)};
        for(i=0;i<len;i+=k)
        {
            w=(cp){1,0};
            for(j=0;j<(k>>1);j++,w=w*wn)
            {
                t=w*s[i+j+(k>>1)];
                s[i+j+(k>>1)]=s[i+j]-t;
                s[i+j]=s[i+j]+t;
            }
        }
    }
}
void FFT_Main(int len)
{
    FFT(a,len,1); FFT(b,len,1);
    for(int i=0;i<len;i++) c[i]=a[i]*b[i];
    FFT(c,len,-1);
    for(int i=0;i<len;i++) c[i].x/=len;
}

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

int main()
{
    scanf("%d",&T);
    while(T--) {
    
    memset(v,0,sizeof(v)); memset(s,0,sizeof(s)); memset(r,0,sizeof(r)); memset(ans,0,sizeof(ans));
    int i,j,maxn=0,len,L,num;
    scanf("%d",&n);
    for(i=1;i<=n;i++) v[i]=gint(), s[i]=s[i-1]+v[i], maxn=max(maxn,s[i]);
    if(!maxn)
    {
        for(i=1;i<=n;i++) 
            ans[0]+=1ll*(i+1)*i/2;
        printf("%lld\n",ans[0]);
        continue;
    }
    
    for(len=1,L=0;len<maxn+maxn+1;len<<=1,L++);
    for(i=0;i<len;i++) r[i]=(r[i>>1]>>1)|((i&1)<<(L-1));
    
    memset(a,0,sizeof(a)); memset(b,0,sizeof(b));
    for(i=0;i<=n;i++) a[s[i]].x+=i;
    for(i=0;i<=n;i++) b[-s[i]+maxn].x++;
    FFT_Main(len);
    for(i=maxn;i<=(maxn<<1);i++) ans[i-maxn]=(ll)(c[i].x+0.5);
    
    memset(a,0,sizeof(a)); memset(b,0,sizeof(b));
    for(i=0;i<=n;i++) a[s[i]].x++;
    for(i=0;i<=n;i++) b[-s[i]+maxn].x+=i;
    FFT_Main(len);
    for(i=maxn;i<=(maxn<<1);i++) ans[i-maxn]-=(ll)(c[i].x+0.5);
    
    for(i=1,num=0;i<=n;i++) 
        if(!v[i]){ num++; ans[0]+=1ll*(num+1)*(num)/2ll; }
        else{ num=0; }
    for(i=0;i<=maxn;i++) printf("%lld\n",ans[i]);
    //puts("");    
    }
    return 0;

}