FFT+生成函数--hdu5307 He is Flying

传送门
题意：给一个数列，若有一个数对 $(i,j)$ 满足 $sum[i]-sum[j-1]==S$ ，则得到 $i-(j-1)$ 的收益，求 $S$ 取 $0$ 到 $sum_n$ 的每一个值时，各自的全部收益。

看那个范围就是典型的 $FFT$ 啊···
不过这个题的构造还是很巧妙的，因为要让 $sum[i]-sum[j-1]==S$ ，所以这就可看做 $FFT$ 中的定值部分，然后收益怎么求呢？

可以想到题目大概就是要构造一个生成函数让指数部分为 $sum$ 值，系数就是要求的答案，那么可以这样构造：
$\sum_i\sum_j (ix^{sum_i})*x^{-sum_{j-1}}-x^{sum_i}*(j-1)x^{-sum_{j-1}}$
发现左边和右边乘起来指数都是 $sum[i]-sum[j-1]$ ，两个系数相减就是 $i-(j-1)$ ，于是就可以用 $FFT$ 做了

但是想吐槽 $hdu$ 有剧毒，一开始数组开小了显示 $tle$ ，后来精度不够开了 $long\ double$ 才能过，不过也可以写 $NTT$ ，但我懒得写了

哦对
$0$ 的部分要单独算，因为有 $0$ 的话就可以用两边的相减，再乘上区间长度加起来，用 $FFT$ 比较麻烦，输入的时候就可以顺便算了

代码如下：

#include<iostream>
#include<cstdio>
#include<algorithm>
#include<cstring>
#include<cmath>
#define maxm 200005
#define maxn 100005
#define LL long long
using namespace std;
const long double Pi=acos(-1.0);

inline int rd(){
	int x=0,f=1;char c=getchar();
	while(c<'0' || c>'9') f=c=='-'?-1:1,c=getchar();
	while(c<='9' && c>='0') x=x*10+c-'0',c=getchar();
	return x*f;
} 

struct complex{
	long double x,y;
	complex(long double xx=0,long double yy=0){x=xx,y=yy;}
}a[maxm],b[maxm],c[maxm];
complex operator +(complex a,complex b){return complex(a.x+b.x,a.y+b.y);}
complex operator -(complex a,complex b){return complex(a.x-b.x,a.y-b.y);}
complex operator *(complex a,complex b){return complex(a.x*b.x-a.y*b.y,a.x*b.y+a.y*b.x);}

int t,n,limit,l,rev[maxm],sum[maxn];
LL ans0;

inline void FFT(complex *F,int type){
	for(int i=0;i<limit;i++)
		if(i<rev[i]) swap(F[i],F[rev[i]]);
	for(int mid=1;mid<limit;mid<<=1){
		complex Wn(cos(Pi/mid),1.0*type*sin(Pi/mid));
		for(int r=mid<<1,j=0;j<limit;j+=r){
			complex w(1,0);
			for(int k=0;k<mid;k++,w=w*Wn){
				complex x=F[j+k],y=w*F[j+mid+k];
				F[j+k]=x+y,F[j+mid+k]=x-y;
			}
		}
	}
	if(type==-1)
		for(int i=0;i<limit;i++) F[i].x/=limit;
}

int main(){
	t=rd();
	while(t--){
		n=rd(); ans0=0; int cnt=0;
		for(int i=1;i<=n;i++){
			int x=rd();
			sum[i]=sum[i-1]+x;
			if(!x){
				++cnt;
				ans0+=1LL*cnt*(cnt+1)/2;
			}
			else cnt=0;
		}
		limit=1,l=0;
		while(limit<=2*sum[n]) limit<<=1,++l;
		for(int i=0;i<=limit;i++)
			rev[i]=(rev[i>>1]>>1)|((i&1)<<(l-1)),
			a[i].x=a[i].y=b[i].x=b[i].y=0;
		for(int i=1;i<=n;i++)
			a[sum[i]].x+=i,b[sum[n]-sum[i-1]].x+=1;
		FFT(a,1); FFT(b,1);
		for(int i=0;i<=limit;i++) {
			c[i]=a[i]*b[i];
			a[i].x=a[i].y=b[i].x=b[i].y=0;
		}
		FFT(c,-1);
		for(int i=1;i<=n;i++)
			a[sum[i]].x+=1,b[sum[n]-sum[i-1]].x+=i-1;
		FFT(a,1); FFT(b,1);
		for(int i=0;i<=limit;i++) a[i]=a[i]*b[i];
		FFT(a,-1);
		for(int i=0;i<=limit;i++) c[i]=c[i]-a[i];
		printf("%lld\n",ans0);
		for(int i=1;i<=sum[n];i++) printf("%lld\n",(LL)(c[i+sum[n]].x+0.5));
	}
	return 0;
}

FFT+生成函数--hdu5307 He is Flying

猜你喜欢