[BZOJ3771]Triple（FFT+生成函数）

=== ===

这里放传送门

=== ===

题解

这个题要求所有组合的方案数，并且还要求具体组合出来的方案，可以想到利用生成函数。因为生成函数乘出来以后它的指数上就记录了所有方案，那么首先构造出只有一个东西的生成函数，也就是说如果有一个价值为$i$的东西就把$x^i$这一项的系数赋值为1，设这个生成函数为$G(x)$。求出$G^3(x)$就求出了在所有东西里面选三次的方案数，设这个数字为$X_1$。
但是这样选可能存在有一个东西选了两次或者选了三次的方案数，要把它们减去才能求出恰好选了三个不重复的东西的方案数。要求有一个东西选了至少两次的方案数，就在生成函数里强行让它选两次，也就是如果有一个价值为$i$的东西就把$x^{2i}$这一项的系数赋值为1，设这个生成函数为$G_2(x)$，求出$G_2(x) \times G(x)$就是有一个东西至少选了两次的方案数，设为$X_2$。
用组合数可以得到$X_2$中的每一种方案在$X_1$中都被计算了三次，所以要控制一下容斥系数。并且这样处理完以后同一个东西选了三次的方案还被多减了两次所以要加回来。
这样计算完了以后就得到了恰好选三个不重复东西的方案数。用类似的方法可以得到恰好选两个不重复东西的方案数，选一个的直接算就可以了。
所有的生成函数卷积都可以用FFT进行优化。

代码

#include<cmath>
#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
const double PI=4*atan(1);
int n,v[40010],N,M,L,R[300010],f[300010],g[300010],ans[300010],Max;
struct Complex{
    double x,y;
    Complex(double X=0,double Y=0){x=X;y=Y;}
    Complex operator + (const Complex &a){return Complex(x+a.x,y+a.y);}
    Complex operator - (const Complex &a){return Complex(x-a.x,y-a.y);}
    Complex operator * (const Complex &a){return Complex(x*a.x-y*a.y,x*a.y+y*a.x);}
}a[300010],b[300010];
void FFT(Complex *a,int N,int opt){
    Complex wn,w,x,y;
    for (int i=0;i<N;i++)
      if (i<R[i]) swap(a[i],a[R[i]]);
    for (int k=1;k<N;k<<=1){
        wn=Complex(cos(PI/k),opt*sin(PI/k));
        for (int p=(k<<1),i=0;i<N;i+=p){
            w=Complex(1,0);
            for (int j=0;j<k;j++){
                x=a[i+j];y=w*a[i+j+k];
                a[i+j]=x+y;a[i+j+k]=x-y;
                w=w*wn;
            }
        }
    }
}
int main()
{
    scanf("%d",&n);
    for (int i=1;i<=n;i++){
        scanf("%d",&v[i]);
        a[v[i]].x=1;b[2*v[i]].x=1;
        Max=max(Max,v[i]);
    }
    M=3*Max;
    for (N=1;N<=M;N<<=1) ++L;
    for (int i=0;i<=N;i++)
      R[i]=(R[i>>1]>>1)|((i&1)<<(L-1));
    FFT(a,N,1);FFT(b,N,1);
    for (int i=0;i<=N;i++) b[i]=b[i]*a[i];
    for (int i=0;i<=N;i++) a[i]=a[i]*a[i]*a[i];
    FFT(a,N,-1);FFT(b,N,-1);
    for (int i=0;i<=N;i++) f[i]=(int)(a[i].x/N+0.5);//f[i]是选了三个东西以后的 
    for (int i=0;i<=N;i++) g[i]=(int)(b[i].x/N+0.5);//g[i]是先选了一个再选了两个的 
    for (int i=0;i<=N;i++) f[i]-=3*g[i];
    for (int i=1;i<=n;i++) f[3*v[i]]+=2;
    for (int i=0;i<=N;i++) ans[i]=f[i]/6;
    memset(a,0,sizeof(a));
    for (int i=1;i<=n;i++) a[v[i]].x=1;
    FFT(a,N,1);
    for (int i=0;i<=N;i++) a[i]=a[i]*a[i];
    FFT(a,N,-1);
    for (int i=0;i<=N;i++) f[i]=(int)(a[i].x/N+0.5);
    for (int i=1;i<=n;i++) f[2*v[i]]--;//f[i]是只选两个东西的 
    for (int i=0;i<=N;i++) ans[i]+=f[i]/2;
    for (int i=1;i<=n;i++) ans[v[i]]++;//最后加上只选一个东西的 
    for (int i=1;i<=N;i++)
      if (ans[i]!=0) printf("%d %d\n",i,ans[i]);
    return 0;
}