[BZOJ 3771] Triple(FFT+生成函数)
题面
给出 n个物品，价值为别为(w_i)且各不相同，现在可以取1个、2个或3个，问每种价值和有几种情况？
分析
这种计数问题容易想到生成函数。
设生成函数(A(x)=sum_{i=1}^{n} x^{w_i})，指数为价值，系数为选的方案数。A表示每种物品取1个的方案数。同理，我们可以写出每种物品取2个和3个的生成函数。
(B(x)=sum_{i=1}^{n} x^{2w_i})
(C(x)=sum_{i=1}^{n} x^{3w_i})
然后就开始大力容斥.
取3个不同物品的情况
直接取3个物品的方案数为(A^3(x))，但是我们还需要减去重复的,如((a,a,b),(a,b,a))就算同一种情况。选2个物品(a)的方案为(B(x))，再选一个物品(b)的方案为(A(x)),任意排列有3种。因此要减(3A(x)B(x))
然而每种物品取3个((a,a,a))这样的方案会被减去3次，而实际上只需要减去1次，所以还要加回(2C(x))
注意到((a,b,c))的6种不同排列方案只算一次。总答案为
[frac{A^3(x)-3A(x)B(x)+2C(x)}{6}
]
取2个不同物品的情况
直接取3个物品的方案为(A^2(x))。重复的((a,a))这种情况的方案为(B(x)),并且((a,b))的2种排列只算1次。总答案为
[frac{A^2(x)-B(x))}{2}
]
取1个不同物品的方案
很简单，就是(A(x))
综上，总答案为
[frac{A^3(x)-3A(x)B(x)+2C(x)}{6}+frac{A^2(x)-B(x))}{2}+A(x)
]
先把A,B,C用FFT转成点值表达式然后相乘，再逆变换一下就得到答案.
代码
#include<iostream>
#include<cstdio>
#include<cstring>
#include<cmath>
#define maxn 400000
using namespace std;
typedef long long ll;
const double pi=acos(-1.0);
struct com{
double real;
double imag;
com(){
}
com(double _real,double _imag){
real=_real;
imag=_imag;
}
com(double x){
real=x;
imag=0;
}
void operator = (const com x){
this->real=x.real;
this->imag=x.imag;
}
void operator = (const double x){
this->real=x;
this->imag=0;
}
friend com operator + (com p,com q){
return com(p.real+q.real,p.imag+q.imag);
}
friend com operator + (com p,double q){
return com(p.real+q,p.imag);
}
friend com operator - (com p,com q){
return com(p.real-q.real,p.imag-q.imag);
}
friend com operator - (com p,double q){
return com(p.real-q,p.imag);
}
friend com operator * (com p,com q){
return com(p.real*q.real-p.imag*q.imag,p.real*q.imag+p.imag*q.real);
}
friend com operator * (com p,double q){
return com(p.real*q,p.imag*q);
}
friend com operator / (com p,double q){
return com(p.real/q,p.imag/q);
}
void print(){
printf("%lf + %lf i ",real,imag);
}
};
void fft(com *x,int n,int type){
static int rev[maxn+5];
int tn=1,k=0;
while(tn<n){
k++;
tn*=2;
}
for(int i=0;i<n;i++) rev[i]=(rev[i>>1]>>1)|((i&1)<<(k-1));
for(int i=0;i<n;i++) if(i<rev[i]) swap(x[i],x[rev[i]]);
for(int len=1;len<n;len*=2){
int sz=len*2;
com wn1=com(cos(2*pi/sz),type*sin(2*pi/sz));
for(int l=0;l<n;l+=sz){
int r=l+len-1;
com wnk=1;
for(int i=l;i<=r;i++){
com tmp=x[i+len];
x[i+len]=x[i]-wnk*tmp;
x[i]=x[i]+wnk*tmp;
wnk=wnk*wn1;
}
}
}
if(type==-1){
for(int i=0;i<n;i++) x[i].real/=n;
}
}
void mul(com *a,com *b,com *ans,int n){
// fft(a,n,1);
// fft(b,n,1);
//避免多次fft
for(int i=0;i<n;i++) ans[i]=a[i]*b[i];
fft(ans,n,-1);
}
int n;
int val[maxn+5];
com a[maxn+5],b[maxn+5],c[maxn+5];
com ans[maxn+5];
int main(){
scanf("%d",&n);
int mv=0;
for(int i=1;i<=n;i++){
scanf("%d",&val[i]);
a[val[i]]=a[val[i]]+1;
b[val[i]*2]=b[val[i]*2]+1;
c[val[i]*3]=c[val[i]*3]+1;
mv=max(mv,val[i]);
}
int tn=1,k=0;
while(tn<mv*3){
k++;
tn*=2;
}
fft(a,tn,1);
fft(b,tn,1);
fft(c,tn,1);
for(int i=0;i<tn;i++){
ans[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];
}
fft(ans,tn,-1);
for(int i=0;i<=mv*3;i++){
if(ll(ans[i].real+0.5)){
printf("%d %lld
",i,ll(ans[i].real+0.5));
}
}
}