[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));
        }
    }
}