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

【题意】给定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=300010;
const double PI=acos(-1);
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=0;k<n;k++)o[k]=complex<double>(cos(2*PI*k/n),sin(2*PI*k/n)),oi[k]=conj(o[k]);
    }
    void transform(complex<double>*a,int n,complex<double>*o){
        int k=0;
        while((1<<k)<n)k++;
        for(int i=0;i<n;i++){
            int t=0;
            for(int j=0;j<k;j++)if(i&(1<<j))t|=(1<<(k-j-1));
            if(i<t)swap(a[i],a[t]);
        }
        for(int l=2;l<=n;l*=2){
            int m=l/2;
            for(complex<double>*p=a;p!=a+n;p+=l){
                for(int i=0;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=0;i<n;i++)a[i]/=n;}
}
int main(){
    scanf("%d",&n);
    int mx=0,t;
    for(int i=0;i<n;i++){
        scanf("%d",&t);
        f[t].real(1);g[t*2].real(1),h[t*3].real(1);
        mx=max(mx,t*3);
    }
    n=1;
    while(n<2*mx+2)n*=2;
    fft::init(n);
    fft::dft(f,n);fft::dft(g,n);fft::dft(h,n);
    complex<double>tmp2=(2),tmp3=(3),tmp6=(6);
    for(int i=0;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=0;i<n;i++)ans[i]=(int)(floor(f[i].real()+0.5));
    for(int i=0;i<n;i++)if(ans[i])printf("%d %d\n",i,ans[i]);
    return 0;
}