HDU5307 He is Flying（FFT）

HDU5307 He is Flying

原題地址：http://acm.hdu.edu.cn/showproblem.php?pid=5307

題意：
給定n個非負整數a1...an$a_1...a_n$ ,對所有0≤S≤∑ai$0\le S\le \sum a_i$ ，輸出所有ai$a_i$ 和爲S的區間的長度和.

數據範圍
1≤n≤105$1\le n\le {10}^5$ ，∑ai≤5×104.$\sum a_i\le 5\times {10}^4.$

題解：

令si$s_i$ 爲ai$a_i$ 的前綴和，那麼對於區間[i,j]，對si−sj−1$s_i-s_{j-1}$ 有着j−i+1$j-i+1$ 的貢獻。
易於想到的是xsi∗x−sj−1=xsi−sj−1$x^{s_i}\ast x^{-s_{j-1}}=x^{s_i-s_{j-1}}$ ，
由於每個區間的貢獻是長度i−j+1$i-j+1$ ，
可以構造出多項式∑xsii ×xsi−1−∑xsi ×xsi−1(i−1)$\sum x^{s_i}i\times x^{s_{i-1}}-\sum x^{s_i}\times x^{s_{i-1}}(i-1)$
xs$x^s$ 的係數就是s的答案。
用FFT計算即可。

1A開心。

代碼：

#include<cstdio>
#include<iostream>
#include<cstring>
#include<algorithm>
#include<cmath>
#define LL long long
#define LD long double
const long double Pi=acos(-1);
using namespace std;
const int N=200005;
inline int read()
{
    int ret=0,w=1; char ch=getchar();
    while((ch<'0'||ch>'9')&&ch!='-') ch=getchar();
    if(ch=='-') w=-1,ch=getchar();
    while(ch>='0'&&ch<='9') ret=(ret<<1)+(ret<<3)+ch-'0',ch=getchar();
    return ret*w;
}
int T,n,s[N],R[N],p,len;
LL ans[N];
struct Virt
{
    LD r,i;
    Virt(){}
    Virt(LD r,LD i):r(r),i(i){}
    Virt operator+(const Virt &A){return Virt(r+A.r,i+A.i);}
    Virt operator-(const Virt &A){return Virt(r-A.r,i-A.i);}
    Virt operator*(const Virt &A){return Virt(r*A.r-i*A.i,r*A.i+i*A.r);}
}a[N],b[N],omg[N],_omg[N];
void FFT(Virt *x,int opt)
{
    for(int i=0;i<len;i++) if(i<R[i]) swap(x[i],x[R[i]]);
    Virt *w; if(opt==1) w=omg; else w=_omg;
    for(int m=2;m<=len;m<<=1)
    {
        int l=m>>1;
        for(int j=0;j<len;j+=m)
        for(int i=0;i<l;i++)
        {
            Virt y=w[len/m*i]*x[j+i+l];
            x[j+i+l]=x[j+i]-y;
            x[j+i]=x[j+i]+y;
        }
    }
    if(opt==-1) for(int i=0;i<len;i++) x[i].r/=len;
}
int main()
{
    T=read();
    while(T--)
    {
        n=read(); LL ret=0; LL cnt=0; s[0]=0;
        for(int i=1;i<=n;i++) 
        {
            s[i]=read();
            if(s[i]==0) {cnt++,ret+=1LL*(1LL+cnt)*cnt/2LL;}
            else cnt=0;
            s[i]+=s[i-1];
        }
        printf("%lld\n",ret); 
        for(p=0,len=1;len<=2*s[n];len<<=1,p++);
        for(int i=0;i<len;i++){omg[i]=Virt(cos(2.0*Pi/(LD)len*i),sin(2.0*Pi/(LD)len*i)); _omg[i]=Virt(omg[i].r,-omg[i].i);}
        R[0]=0; for(int i=1;i<len;i++) R[i]=(R[i>>1]>>1)|((i&1)<<(p-1));
        for(int i=0;i<=len;i++) a[i]=b[i]=Virt(0,0);
        for(int i=1;i<=n;i++) {a[s[i]].r+=(LD)i; b[-s[i-1]+s[n]].r+=(LD)1.0;}
        FFT(a,1); FFT(b,1);
        for(int i=0;i<len;i++) a[i]=a[i]*b[i];
        FFT(a,-1);
        for(int i=1;i<=s[n];i++) ans[i]=(LL)(a[i+s[n]].r+0.5);
        for(int i=0;i<=len;i++) a[i]=b[i]=Virt(0,0);
        for(int i=1;i<=n;i++) {a[s[i]].r+=1.0; b[-s[i-1]+s[n]].r+=(LD)(i-1);}
        FFT(a,1); FFT(b,1);
        for(int i=0;i<len;i++) a[i]=a[i]*b[i];
        FFT(a,-1);
        for(int i=1;i<=s[n];i++) ans[i]-=(LL)(a[i+s[n]].r+0.5);
        for(int i=1;i<=s[n];i++) printf("%lld\n",ans[i]);
    }   
    return 0;
}