【BZOJ3771】Triple（FFT水题）

点此看题面

大致题意： 你可以在\(n\)个数中选出\(1\)/\(2\)/\(3\)个数，求所有可能的和以及对应的方案数。

大致思路

我们可以求出生成函数\(F1(x)=\sum_{i=0}^{\infty}p_ix^i\)，其中\(p_i\)表示\(i\)的个数，容易发现这也就是选一个数和为\(i\)的方案数。

考虑选出两个数可以把原式平方得到\(F1(x)^2\)，由于一个数不能选多次，所以我们要将其减去\(G2(x)=\sum_{i=0}^{\infty}p_ix^{2i}\)。

然后发现一组\((i,j)\)会被计算两次，所以我们再将系数除以\(2\)，也就是：

\[F2(x)=\frac{F1(x)^2-G2(x)}2 \]

同理，我们把\(F1(x)\)和\(F2(x)\)卷起来得到\(F1(x)F2(x)\)，把\(F1(x)\)和\(G2(x)\)卷起来并减去\(\sum_{i=0}^{\infty}p_ix^{3i}\)得到\(G3(x)\)（因为\(F2(x)\)保证无重复，故\(F1(x)F2(x)\)不会出现一个数选三次的情况，但\(F1(x)G2(x)\)卷起来会有这种情况，不减去就会导致最终答案变小）。

此时一组\((i,j,k)\)会被计算\(3\)次，所以我们再将系数除以\(3\)，也就是：

\[F3(x)=\frac{F1(x)F2(x)-G3(x)}3 \]

最终和为\(i\)的方案数就是\(F1(x),F2(x),F3(x)\)中\(i\)次项系数之和。

具体实现可以详见代码。

代码

#include<bits/stdc++.h>
#define Tp template<typename Ty>
#define Ts template<typename Ty,typename... Ar>
#define Reg register
#define RI Reg int
#define Con const
#define CI Con int&
#define I inline
#define W while
#define N 40000
#define LL long long
#define DB double
using namespace std;
int n;LL f1[3*N+5],f2[3*N+5],g2[3*N+5],f3[3*N+5],g3[3*N+5];
class FastIO
{
	private:
		#define FS 100000
		#define tc() (A==B&&(B=(A=FI)+fread(FI,1,FS,stdin),A==B)?EOF:*A++)
		#define pc(c) (C==E&&(clear(),0),*C++=c)
		#define D isdigit(c=tc())
		int T;char c,*A,*B,*C,*E,FI[FS],FO[FS],S[FS];
	public:
		I FastIO() {A=B=FI,C=FO,E=FO+FS;}
		Tp I void read(Ty& x) {x=0;W(!D);W(x=(x<<3)+(x<<1)+(c&15),D);}
		Tp I void write(Ty x) {W(S[++T]=x%10+48,x/=10);W(T) pc(S[T--]);}
		Tp I void write(Con Ty& x,Con char& y) {write(x),pc(y);}
		I void clear() {fwrite(FO,1,C-FO,stdout),C=FO;}
}F;
class Poly//多项式乘法板子
{
	private:
		int P,L,R[12*N+5];DB Pi;
		struct node
		{
			DB x,y;I node(Con DB& a=0,Con DB& b=0):x(a),y(b){}
			I node operator + (Con node& o) Con {return node(x+o.x,y+o.y);}
			I node operator - (Con node& o) Con {return node(x-o.x,y-o.y);}
			I node operator * (Con node& o) Con {return node(x*o.x-y*o.y,x*o.y+y*o.x);}
		}A[12*N+5],B[12*N+5];
		I void T(node *s,CI op)
		{
			RI i,j,k;node x,y,U,S;for(i=0;i^P;++i) i<R[i]&&(x=s[i],s[i]=s[R[i]],s[R[i]]=x,0);
			for(i=1;i^P;i<<=1) for(U=node(cos(Pi/i),op*sin(Pi/i)),j=0;j^P;j+=i<<1)
				for(S=1,k=0;k^i;++k,S=S*U) s[j+k]=(x=s[j+k])+(y=S*s[i+j+k]),s[i+j+k]=x-y;
		}
	public:
		I Poly() {Pi=acos(-1);}
		I void FFT(CI n,LL *a,CI m,LL *b,LL *c)
		{
			RI i;P=1,L=0;W(P<=n+m) P<<=1,++L;
			for(i=0;i^P;++i) R[i]=(R[i>>1]>>1)|((i&1)<<L-1),A[i]=B[i]=0;
			for(i=0;i<=n;++i) A[i]=a[i];for(i=0;i<=m;++i) B[i]=b[i];
			for(T(A,1),T(B,1),i=0;i^P;++i) A[i]=A[i]*B[i];
			for(T(A,-1),i=0;i<=n+m;++i) c[i]=A[i].x/P+0.5;
		}
}P;
int main()
{
	RI i,x;for(F.read(n),i=1;i<=n;++i) F.read(x),++f1[x];//F1(x)
	for(P.FFT(N,f1,N,f1,f2),i=0;i<=N;++i) g2[i<<1]=f1[i];//F1(x)^2和G2(x)
	for(i=0;i<=2*N;++i) f2[i]=(f2[i]-g2[i])/2;//F2(x)
	for(P.FFT(N,f1,N<<1,f2,f3),P.FFT(N,f1,N<<1,g2,g3),i=0;i<=N;++i) g3[3*i]-=f1[i];//F1(x)F2(x)和G3(x)
	for(i=0;i<=3*N;++i) f3[i]=(f3[i]-g3[i])/3;//F3(x)
	for(i=0;i<=3*N;++i) f1[i]+f2[i]+f3[i]&&//如果可能得到这个和
		(F.write(i,' '),F.write(f1[i]+f2[i]+f3[i],'\n'),0);//输出和以及对应方案数
	return F.clear(),0;
}