HDU 5307 He is Flying ——FFT

卷积的妙用，显然我们可以求出所有符合条件的右端点的和，然后减去左端点的和。

就是最后的答案。然后做一次前缀和，然后就变成了统计差是一个定值的情况。

令$A(s[i])++$ $B(s[i])+=i$

然后卷积一次就可以了，然后用后半部分减去前半部分即可。

并不需要两次FFT

然后发现$0$的情况会导致重叠。所以特判就好了。

#include <map>
#include <cmath>
#include <queue>
#include <cstdio>
#include <cstring>
#include <iostream>
#include <algorithm>
using namespace std;
#define F(i,j,k) for (int i=j;i<=k;++i)
#define D(i,j,k) for (int i=j;i>=k;--i)
#define ll long long
#define mp make_pair
#define double long double
#define maxn 400005

struct Complex{
    double x,y;
    Complex () {}
    Complex (double _x,double _y) {x=_x;y=_y;}
    Complex operator + (Complex a) {return Complex(x+a.x,y+a.y);}
    Complex operator - (Complex a) {return Complex(x-a.x,y-a.y);}
    Complex operator * (Complex a) {return Complex(x*a.x-y*a.y,x*a.y+y*a.x);}
    void init(){x=0;y=0;}
}A[maxn],B[maxn],C[maxn],ans[maxn];

int t,n,a[maxn],sum,rev[maxn],m,len,top;
const double pi=acos(-1.0);

void FFT(Complex * x,int n,int flag)
{
    F(i,0,n-1) if (rev[i]>i) swap(x[rev[i]],x[i]);
    for (int m=2;m<=n;m<<=1)
    {
        Complex wn=Complex(cos(2*pi/m),flag*sin(2*pi/m));
        for (int i=0;i<n;i+=m)
        {
            Complex w=Complex(1.0,0);
            for (int j=0;j<(m>>1);++j)
            {
                Complex u=x[i+j],v=x[i+j+(m>>1)]*w;
                x[i+j]=u+v;x[i+j+(m>>1)]=u-v;
                w=w*wn;
            }
        }
    }
}

int main()
{
    scanf("%d",&t);
    while (t--)
    {
        scanf("%d",&n);int flag=1,cnt=0;ll ans0=0;
        F(i,1,n){scanf("%d",&a[i]);if(a[i]==0&&flag)cnt++;else{F(i,1,cnt) ans0+=(ll)i*(cnt-i+1);cnt=0;}a[i]+=a[i-1];}
        F(i,1,cnt) ans0+=(ll)i*(cnt-i+1);cnt=0;
        sum=a[n];top=n;a[0]=0;n=2*(sum+1)+3;m=1;len=0;while (m<=n) m<<=1,len++;n=m;
        F(i,0,n-1){int ret=0,t=i;F(j,1,len)ret<<=1,ret|=t&1,t>>=1;rev[i]=ret;}
        F(i,0,n-1) A[i].init(),B[i].init(),ans[i].init();F(i,0,top) A[sum-a[i]].x+=1.0,B[a[i]].x+=1.0*i;
        FFT(A,n,1);FFT(B,n,1);F(i,0,n-1)C[i]=A[i]*B[i];FFT(C,n,-1);F(i,0,n-1)C[i].x=C[i].x/n;
        F(i,0,sum) ans[i].x+=C[sum+i].x-C[sum-i].x;
        printf("%lld\n",ans0);F(i,1,sum) printf("%lld\n",(ll)(ans[i].x+0.5));
    }
}