【BZOJ】3771: Triple FTT+生成函数

时间:2022-05-06 06:25:34

【题意】给定n个物品,价值为$a_i$,物品价格互不相同,求选一个或两个或三个的价值为x的方案数,输出所有存在的x和对应方案数。$ai<=40000$。

【算法】生成函数+FFT

【题解】要求价值为x的方案数,就定义价值为“大小”(即指数),方案数为“元素个数”(即系数),物品为“选择项”(即多项式)。

设$f(x)$表示选择一个物品,价值为x的方案数。

虽然很容易想到$f^3(x)$,但因为不能重复选物品所以必须去重。而且不容易在$f^3(x)$中表示出选一个或两个物品。

①选择一个物品的方案数:$f$

考虑$f^2$表示选择两个物品价值和为x的可重排列数

去重需要减去两个物品相同的情况,设$g(x)$表示选择两个相同物品价值和为x的方案数,显然g可以直接得到。(枚举每个物品,在价值*2的系数处+1)。

因为生成函数乘积在两个物品相同的时候实际上没有排列,所以要先减再去排列。

②选择两个物品的方案数:$\frac{f^2-g}{2}$

最后三个物品同理,设$h(x)$表示选择三个相同物品价值和为x的方案数,需要排除BAA,ABA,AAB,AAA的情况。

其中BAA,ABA,AAB相当于选择两个相同物品后再选一个物品(无论是否再相同),即$f*g$。

但这样会把AA重复减去三次,实际上只需要减去一次,所以容斥加回,即$h$。

③选择三个物品的方案数:$\frac{f^3-3*f*g+2*h}{2}$

复杂度O(n log n)。

注意:先将f,g,h进行DFT,全部计算答案后再进行IDFT,才能保证精度。

#include<cstdio>
#include<cstring>
#include<algorithm>
#include<complex>
#include<cmath>
using namespace std;
const int maxn=;
const double PI=acos(-);
int n,ans[maxn];
complex<double>f[maxn],g[maxn],h[maxn];
namespace fft{
complex<double>o[maxn],oi[maxn];
void init(int n){
for(int k=;k<n;k++)o[k]=complex<double>(cos(*PI*k/n),sin(*PI*k/n)),oi[k]=conj(o[k]);
}
void transform(complex<double>*a,int n,complex<double>*o){
int k=;
while((<<k)<n)k++;
for(int i=;i<n;i++){
int t=;
for(int j=;j<k;j++)if(i&(<<j))t|=(<<(k-j-));
if(i<t)swap(a[i],a[t]);
}
for(int l=;l<=n;l*=){
int m=l/;
for(complex<double>*p=a;p!=a+n;p+=l){
for(int i=;i<m;i++){
complex<double>t=p[i+m]*o[n/l*i];
p[i+m]=p[i]-t;
p[i]+=t;
}
}
}
}
void dft(complex<double>*a,int n){transform(a,n,o);}
void idft(complex<double>*a,int n){transform(a,n,oi);for(int i=;i<n;i++)a[i]/=n;}
}
int main(){
scanf("%d",&n);
int mx=,t;
for(int i=;i<n;i++){
scanf("%d",&t);
f[t].real();g[t*].real(),h[t*].real();
mx=max(mx,t*);
}
n=;
while(n<*mx+)n*=;
fft::init(n);
fft::dft(f,n);fft::dft(g,n);fft::dft(h,n);
complex<double>tmp2=(),tmp3=(),tmp6=();
for(int i=;i<n;i++)f[i]=(f[i]*f[i]*f[i]-tmp3*f[i]*g[i]+tmp2*h[i])/tmp6+(f[i]*f[i]-g[i])/tmp2+f[i];
fft::idft(f,n);
for(int i=;i<n;i++)ans[i]=(int)(floor(f[i].real()+0.5));
for(int i=;i<n;i++)if(ans[i])printf("%d %d\n",i,ans[i]);
return ;
}