【BZOJ3771】Triple

剛學的FFT。。證明好玄乎啊

根據mjs大佬的原話，FFT這種東西不需要理解，背了模板就好

先貼題

鏈接：BZOJ3771 Triple

題意：從n個數中選出1,2或3個數求和，詢問組成每個和的方案數。

思路：生成函數+FFT+容斥原理

假設可選的數表爲 {1,3,4,4,4,8,8,10}$\{1,3,4,4,4,8,8,10\}$ ，那麼構造多項式 A=x+x3+3x4+2x8+x10$A=x+x^3+3x^4+2x^8+x^{10}$ 。怎麼解釋呢？將 A$A$ 表示爲若干個單項式 aixki$a_i{x}^{k_i}$ 之和，其中 ki$k_i$ 爲原本數表中的數， ai$a_i$ 爲 ki$k_i$ 在表中出現的次數。

那答案不就是 A3$A^3$ 各項的係數與次數嗎？

並不是。。將 A3$A^3$ 展開，會發現某項會自乘，從而對答案產生貢獻。

這時候我們就要請出 客廳原理 容斥原理了

再構造多項式 B=x2+x6+3x8+2x16+x20$B=x^2+x^6+3x^8+2x^{16}+x^{20}$ 同 A$A$ 的原理，但 ki$k_i$ 爲原本數表中數的兩倍， C$C$ 中爲三倍。
這樣 B$B$ 與 C$C$ 即可表示將某個數限制選兩次（三次）的方案數了

於是由容斥原理我們得到

S=A3−3AB+2C6+A2−B2+A$$S=\frac{A^3-3AB+2C}{6}+\frac{A^2-B}{2}+A$$

S$S$ 爲最終的答案

於是我們就可以用FFT在 O(nlogn)$O(n\log n)$ 的時間內求出多項式 S$S$ 啦

上代碼

#include<iostream>
#include<cstdio>
#include<cmath>

using namespace std;

struct CComplexNumber{
    double r,i;
    CComplexNumber():r(0.),i(0.){}
    CComplexNumber(const double& _r,const double& _i):r(_r),i(_i){}
    friend CComplexNumber operator+(const CComplexNumber& c1,const CComplexNumber& c2){return CComplexNumber(c1.r+c2.r,c1.i+c2.i);}
    friend CComplexNumber operator-(const CComplexNumber& c1,const CComplexNumber& c2){return CComplexNumber(c1.r-c2.r,c1.i-c2.i);}
    friend CComplexNumber operator*(const CComplexNumber& c1,const CComplexNumber& c2){return CComplexNumber(c1.r*c2.r-c1.i*c2.i,c1.r*c2.i+c1.i*c2.r);}
    friend CComplexNumber operator*(const double& x,const CComplexNumber& c){return CComplexNumber(x*c.r,x*c.i);}
    friend CComplexNumber operator/(const CComplexNumber& c,const double& x){return CComplexNumber(c.r/x,c.i/x);}
    CComplexNumber operator*=(const CComplexNumber& c){CComplexNumber c0(r*c.r-i*c.i,r*c.i+i*c.r);r=c0.r;i=c0.i;return *this;}
    CComplexNumber operator/=(const double& x){r/=x;i/=x;return *this;}
};

const double Pi=acos(-1);
CComplexNumber x[360001],y[360001],z[360001];
int len=1,N,n,rev[360001],a,bit;

void fft(CComplexNumber *f,int check){
    for(int i=0;i<len;++i)if(i<rev[i])swap(f[i],f[rev[i]]);
    for(int i=1;i<len;i<<=1){
        CComplexNumber wn(cos(Pi/i),check*sin(Pi/i));
        for(int j=0;j<len;j+=(i<<1)){
            CComplexNumber w(1.,0.);
            for(int k=0;k<i;++k){CComplexNumber x=f[j+k],y=w*f[i+j+k];f[j+k]=x+y,f[i+j+k]=x-y;w*=wn;}
        }
    }
    if(check==-1)for(int i=0;i<len;++i)f[i]/=(double)len;
}

int main(){
    scanf("%d",&n);
    for(int i=1;i<=n;++i)scanf("%d",&a),++x[a].r,++y[2*a].r,++z[3*a].r,N=max(N,a);
    while(len<=6*N+1)len*=2,++bit;
    for(int i=0;i<len;++i)rev[i]=((rev[i>>1]>>1)|((1&i)<<bit-1));
    fft(x,1),fft(y,1),fft(z,1);
    for(int i=0;i<len;++i)x[i]=(x[i]*x[i]*x[i]-3*x[i]*y[i]+2*z[i])/6.+(x[i]*x[i]-y[i])/2.+x[i];
    fft(x,-1);
    for(int i=0;i<len;++i)if(round(x[i].r))printf("%d %d\n",i,(int)round(x[i].r));
}