problem
有 n n n 个物品,每个物品有一个权值 a i a_i ai。现在要在这 n n n 个物品中选出 1 1 1 个、 2 2 2 个或 3 3 3 个,总价值就是选出物品的价值和。对于每个可能的总价值,请你计算出有几种可能的方案。
注:如果选了 a a a 和 b b b, ( a , b ) (a,b) (a,b) 和 ( b , a ) (b,a) (b,a) 视为一种方案。如果选了 a a a 和 b b b 和 c c c, ( a , b , c ) , ( b , c , a ) , ( c , a , b ) , ( c , b , a ) , ( b , a , c ) , ( a , c , b ) (a,b,c),(b,c,a),(c,a,b),(c,b,a),(b,a,c),(a,c,b) (a,b,c),(b,c,a),(c,a,b),(c,b,a),(b,a,c),(a,c,b) 视为一种方案。
数据范围: a i ≤ 40000 a_i\le40000 ai≤40000。
solution
把生成函数写出来,用 A ( x ) , B ( x ) , C ( x ) A(x),B(x),C(x) A(x),B(x),C(x) 表示一个物品选 1 / 2 / 3 1/2/3 1/2/3 个的生成函数:
A ( x ) = ∑ i = 1 n x a i B ( x ) = ∑ i = 1 n x 2 a i C ( x ) = ∑ i = 1 n x 3 a i \begin{aligned} A(x)&=\sum_{i=1}^nx^{a_i}\\ B(x)&=\sum_{i=1}^nx^{2a_i}\\ C(x)&=\sum_{i=1}^nx^{3a_i} \end{aligned} A(x)B(x)C(x)=i=1∑nxai=i=1∑nx2ai=i=1∑nx3ai
那么在 n n n 个物品中选 1 1 1 个的方案为:
A ( x ) A(x) A(x)
选 2 2 2 个的方案为:
A 2 ( x ) − B ( x ) 2 \frac{A^2(x)-B(x)}2 2A2(x)−B(x)
就是因为 A 2 ( x ) A^2(x) A2(x) 可能会有同一件物品选两次的情况,所以要减个 B ( x ) B(x) B(x)。
选 3 3 3 个的方案为:
A 3 ( x ) − 3 A ( x ) B ( x ) + 2 C ( x ) 6 \frac{A^3(x)-3A(x)B(x)+2C(x)}6 6A3(x)−3A(x)B(x)+2C(x)
因此我们 DFT 之后对每一项按照上面所述加起来,再 IDFT 回去就可以了。
时间复杂度 O ( n log n ) O(n\log n) O(nlogn)。
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;
}