[BZOJ 3771] Triple(FFT+容斥原理+生成函数)

[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\n",i,ll(ans[i].real+0.5));
		}
	}
}