hdu5307 He is Flying FFT

来源:互联网 发布:贵阳广电网络营业厅 编辑:程序博客网 时间:2024/06/05 00:36

处理一个前缀和,那么ans[x]=sum[i]sum[j]=xij
然后sum[i]sum[j]可以当成卷积做,不过由于FFT只支持乘法,因此可以先做一遍关于i的,再做一遍关于j的,然后相减。
注意0的情况。。。

#include <bits/stdc++.h>using namespace std;#define N 110000#define ll long long#define ld long doubleconst ld PI=acos(-1);int T,n;int sum[N],cnt;ll ans[N],pre[N];struct cp{    ld x,y;    cp(){}    cp(ld x,ld y):x(x),y(y){}    friend cp operator + (const cp &r1,const cp &r2)    {return cp(r1.x+r2.x,r1.y+r2.y);}    friend cp operator - (const cp &r1,const cp &r2)    {return cp(r1.x-r2.x,r1.y-r2.y);}    friend cp operator * (const cp &r1,const cp &r2)    {return cp(r1.x*r2.x-r1.y*r2.y,r1.x*r2.y+r2.x*r1.y);}}a[N<<2],b[N<<2],c[N<<2];void init(){    memset(a,0,sizeof(a));    memset(b,0,sizeof(b));    memset(c,0,sizeof(c));}void FFT(cp *a,int len,int type){    int t=0;    for(int i=0;i<len;i++)    {        if(i>t)swap(a[i],a[t]);        for(int j=len>>1;(t^=j)<j;j>>=1);    }    for(int i=2;i<=len;i<<=1)    {        cp wn(cos(2*PI*type/i),sin(2*PI*type/i));        for(int j=0;j<len;j+=i)        {            cp t,w(1,0);            for(int k=0;k<(i>>1);k++,w=w*wn)            {                t=w*a[j+k+(i>>1)];                a[j+k+(i>>1)]=a[j+k]-t;                a[j+k]=a[j+k]+t;            }        }    }}void solve(cp *a,cp *b,cp *c,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);}int main(){    //freopen("tt.in","r",stdin);    for(int i=1;i<=100000;i++)        pre[i]=pre[i-1]+(ll)i*(i+1)/2;    scanf("%d",&T);    while(T--)    {        memset(ans,0,sizeof(ans));        scanf("%d",&n);cnt=0;        for(int i=1;i<=n;i++)        {            scanf("%d",&sum[i]);            if(sum[i]==0)cnt++;            else ans[0]+=pre[cnt],cnt=0;            sum[i]+=sum[i-1];        }        ans[0]+=pre[cnt];        int len=1;for(;len<=sum[n]<<1;len<<=1);        init();        for(int i=0;i<=n;i++)            a[sum[i]].x+=i,b[sum[n]-sum[i]].x++;        solve(a,b,c,len);        for(int i=1;i<=sum[n];i++)            ans[i]+=(ll)(c[i+sum[n]].x/len+0.1);        init();        for(int i=0;i<=n;i++)            a[sum[i]].x++,b[sum[n]-sum[i]].x+=i;        solve(a,b,c,len);        for(int i=1;i<=sum[n];i++)            ans[i]-=(ll)(c[i+sum[n]].x/len+0.1);        for(int i=0;i<=sum[n];i++)            printf("%lld\n",ans[i]);    }    return 0;}
0 0