HDU-5307 He is Flying (FFT)

Problem Description
JRY wants to drag racing along a long road. There are n sections on the road, the i -th section has a non-negative integer length si .
JRY will choose some continuous sections to race (at an unbelievable speed), so there are totally n(n+1)2 different ways for him to ride.
If JRY rides across from the i -th section to the j -th section, he would gain j?i+1 pleasure. Now JRY wants to know, if he tries all the ways whose length is s ,
what's the total pleasure he can get. Please be aware that in the problem, the length of one section could be zero, which means that the length is so trivial that we can regard it as 0
Input
The first line of the input is a single integer T (T=5) , indicating the number of testcases.

For each testcase, the first line contains one integer n . The second line contains n non-negative integers,
which mean the length of every section. If we denote the total length of all the sections as s , we can guarantee that 0≤s≤50000 and 1≤n≤100000 .
Output
For each testcase, print s+1 lines. The single number in the i -th line indicates the total pleasure JRY can get if he races all the ways of length i?1 .
Sample Input
2 3 1 2 3 4 0 1 2 3
Sample Output
0 1 1 3 0 2 3 1 3 1 6 0 2 7

解题思路：

首先可以得到一个暴力的做法，就是枚举每一段的起点终点将答案计入数组，

输出就好了，时间复杂度$O(n^3)$，前缀和优化就是$O(n^2)$。

分析上述算法，发现如果固定住答案数组的一位（假设现在就是在求第$x$项的答案，设前$i$项$s[]$的前缀和为$sum[i]$）

那么上述算法可以表述为：

$Ans(x)=\sum\limits_{i=1}^{n}\sum\limits_{j=i}^{n}[sum[j]-sum[i-1]==x](j-i+1)$

由于$sum[j]-sum[i-1]$为固定的，可以认为这是在暗示这题正解是求生成函数第$x$项系数。

那我们就往生成函数上想吧。

那么由于后面是$j-i+1$是加减法关系，由于次数不同，可以认为是两个生成函数分别与全是1的多项式乘积后的累和。

那么考虑是哪两个结果相加。这个好像只能是$j$和$i-1$了吧（指数形式提醒我们）。

设答案的生成函数多项式为$F(x)=f(0)+f(1)x+f(2)x^2+...$

设左右两部分贡献分别为$G$和$H$。（左边是根据$j$列出的多项式，右边是根据$i-1$列出的多项式）

那么

$F=G-H$

由于需要将多项式次数调平，所以:

$G=\sum\limits_{i=1}^{n}(i*x^{sum[i]})\sum\limits_{i=1}^{n}(x^{-sum[i-1]})$

同理：

$H=\sum\limits_{i=1}^{n}((i-1)x^{-sum[i-1]})\sum\limits_{i=1}^{n}x^{sum[i]}$

负数次幂加个大数最后减掉就好了。

注意这个式子并不能处理(i==0)的情况，因为$sum[-1]$并没有被定义。

分类讨论好麻烦，所以直接$O(n)$统计(i==0)的答案即可。

注意一定要写using namespace std;

代码：

#include<cmath>
#include<queue>
#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
const int DLT=50000;
const long double PI=acos(-1.0);
struct cp{
    long double x,y;
    cp(){}
    cp(long double a,long double b){x=a;y=b;}
    cp friend operator + (cp a,cp b){return cp(a.x+b.x,a.y+b.y);}
    cp friend operator - (cp a,cp b){return cp(a.x-b.x,a.y-b.y);}
    cp friend operator * (cp a,cp b){return cp(a.x*b.x-a.y*b.y,a.x*b.y+a.y*b.x);}
}A[200010],B[200010],C[200010],D[200010];
int T;
int n,m;
int lim;
int len;
int s[200010];
int pos[200010];
int sum[200010];
void init(void)
{
    memset(A,0,sizeof(A));
    memset(B,0,sizeof(B));
    memset(C,0,sizeof(C));
    memset(D,0,sizeof(D));
    return ;
}
void FFT(cp *a,double flag)
{
    for(int i=0;i<len;i++)if(i<pos[i])std::swap(a[i],a[pos[i]]);
    for(int i=2;i<=len;i<<=1)
    {
        cp wn(cos(2.00*flag*PI/i),sin(2.00*flag*PI/i));
        for(int j=0;j<len;j+=i)
        {
            cp w(1,0),t;
            for(int k=0;k<(i>>1);k++,w=w*wn)
            {
                t=a[j+k+(i>>1)]*w;
                a[j+k+(i>>1)]=a[j+k]-t;
                a[j+k]=a[j+k]+t;
            }
        }
    }
    return ;
}
void work(void)
{
    init();
    long long a0(0);
    int cnt(0);
    scanf("%d",&n);
    for(int i=1;i<=n;i++)
    {
        scanf("%d",&s[i]);
        if(!s[i])
        {
            cnt++;
            a0+=1ll*cnt*(cnt+1)/2;
        }else cnt=0;
        sum[i]=sum[i-1]+s[i];
    }
    printf("%lld\n",a0);
    while((1<<lim)<(DLT<<1))lim++;len=1<<lim;
    for(int i=0;i<len;i++)pos[i]=(pos[i>>1]>>1)|((i&1)<<(lim-1));
    for(int i=1;i<=n;i++)
    {
        A[sum[i]].x+=i;B[DLT-sum[i-1]].x+=1;
        C[sum[i]].x+=1;D[DLT-sum[i-1]].x+=i-1;
    }
    FFT(A,1),FFT(B,1),FFT(C,1),FFT(D,1);
    for(int i=0;i<len;i++)A[i]=A[i]*B[i]-C[i]*D[i];
    FFT(A,-1);
    for(int i=1;i<=sum[n];i++)printf("%I64d\n",(long long)(A[i+DLT].x/len+0.2));
    return ;
}
int main()
{
    scanf("%d",&T);
    while(T --> 0)work();
    return 0;
}