【BZOJ 3771】Triple

传送门

---

problem

有 $n$ 个物品，每个物品有一个权值 $a_i$。现在要在这 $n$ 个物品中选出 $1$ 个、$2$ 个或 $3$ 个，总价值就是选出物品的价值和。对于每个可能的总价值，请你计算出有几种可能的方案。

注：如果选了 $a$ 和 $b$，$(a,b)$ 和 $(b,a)$ 视为一种方案。如果选了 $a$ 和 $b$ 和 $c$，$(a,b,c),(b,c,a),(c,a,b),(c,b,a),(b,a,c),(a,c,b)$ 视为一种方案。

数据范围：$a_i\le 40000$。

---

solution

把生成函数写出来，用 $A(x),B(x),C(x)$ 表示一个物品选 $1/2/3$ 个的生成函数：

$$\begin{array}{rl}A(x)& =\sum _{i=1}^n{x}^{a_i}\\ B(x)& =\sum _{i=1}^n{x}^{2a_i}\\ C(x)& =\sum _{i=1}^n{x}^{3a_i}\end{array}$$

那么在 $n$ 个物品中选 $1$ 个的方案为：

$$A(x)$$

选 $2$ 个的方案为：

$$\frac{A^2(x)-B(x)}{2}$$

就是因为 $A^2(x)$ 可能会有同一件物品选两次的情况，所以要减个 $B(x)$。

选 $3$ 个的方案为：

$$\frac{A^3(x)-3A(x)B(x)+2C(x)}{6}$$

因此我们 DFT 之后对每一项按照上面所述加起来，再 IDFT 回去就可以了。

时间复杂度 $O(n\log n)$。

---

code

#include<cmath>
#include<cstdio>
#include<cstring>
#include<algorithm>
#define N 1000005
using namespace std;
const double pi=acos(-1);
int n,Max,limit=1,pos[N];
struct complex{
	double x,y;
	complex(){}
	complex(double x,double y):x(x),y(y){}
	friend complex operator+(const complex &a,const complex &b)  {return complex(a.x+b.x,a.y+b.y);}
	friend complex operator-(const complex &a,const complex &b)  {return complex(a.x-b.x,a.y-b.y);}
	friend complex operator*(const complex &a,const complex &b)  {return complex(a.x*b.x-a.y*b.y,a.x*b.y+a.y*b.x);}
	friend complex operator*(double b,const complex &a)  {return complex(a.x*b,a.y*b);}
	friend complex operator/(const complex &a,double b)  {return complex(a.x/b,a.y/b);}
}A[N],B[N],C[N];
void prework(){
	while(limit<=(Max<<2))  limit<<=1;
	for(int i=0;i<limit;++i)
		pos[i]=(pos[i>>1]>>1)|((i&1)*(limit>>1));
}
void DFT(complex f[],int type){
	for(int i=0;i<limit;++i)
		if(pos[i]>i)  swap(f[i],f[pos[i]]);
	for(int mid=1;mid<limit;mid<<=1){
		complex now=complex(cos(pi/mid),type*sin(pi/mid));
		for(int i=0;i<limit;i+=(mid<<1)){
			complex w=complex(1,0);
			for(int j=0;j<mid;++j,w=w*now){
				complex p0=f[i+j],p1=w*f[i+j+mid];
				f[i+j]=p0+p1,f[i+j+mid]=p0-p1;
			}
		}
	}
	if(type==-1)  for(int i=0;i<limit;++i)  f[i].x/=limit;
}
int main(){
	scanf("%d",&n);
	for(int i=1;i<=n;++i){
		int val;
		scanf("%d",&val);
		A[val].x++,B[val*2].x++,C[val*3].x++;
		Max=max(Max,val);
	}
	prework();
	DFT(A,1),DFT(B,1),DFT(C,1);
	for(int i=0;i<limit;++i)  A[i]=(A[i]*A[i]*A[i]-3*A[i]*B[i]+2*C[i])/6+(A[i]*A[i]-B[i])/2+A[i];
	DFT(A,-1);
	for(int i=0;i<limit;++i)
		if(int(A[i].x+0.5))  printf("%d %d\n",i,int(A[i].x+0.5));
	return 0;
}