BZOJ3771：Triple（生成函數+FFT+容斥）

題面
題意：一些數，對於每個x，選3個以下的數，不計順序，和爲x的方案數。

順序就是階乘，先不管，最後乘上就行了

設生成函數爲A
考慮選兩個數，直接卷積後會算重複的，減去兩個相同的方案數
三個的話，會算重選了兩個一樣的
減去後還要加上三個一樣的
可以強行對於每種都選了兩個一樣的，再做一個生成函數B
選三個一樣的，生成函數C

容斥一下發現是大概這樣的
A3−3A∗B+2∗C$A^3-3A\ast B+2\ast C$
FFT加速就行了

#include <iostream>
#include <fstream>
#include <algorithm>
#include <cmath>
#include <ctime>
#include <cstdio>
#include <cstdlib>
#include <cstring>

using namespace std;
#define mmst(a, b) memset(a, b, sizeof(a))
#define mmcp(a, b) memcpy(a, b, sizeof(b))

typedef long long LL;

const double pi=acos(-1);
const int N=144444;

int n,rev[N];

struct yy
{
    double x,y;
    yy(double a=0,double b=0):x(a),y(b){}
};

yy operator +(yy a,yy b){return yy(a.x+b.x,a.y+b.y);}
yy operator -(yy a,yy b){return yy(a.x-b.x,a.y-b.y);}
yy operator *(yy a,yy b){return yy(a.x*b.x-a.y*b.y,a.x*b.y+a.y*b.x);}

void init(int lim)
{
    int k=-1;
    n=1;
    while(n<=lim)
    n<<=1,k++;
    for(int i=0;i<n;i++)
    rev[i]=(rev[i>>1]>>1) | ((i&1)<<k);
}

void fft(yy *a,int ops)
{
    for(int i=0;i<n;i++)
    if(i<rev[i])
    swap(a[i],a[rev[i]]);

    for(int m=1,l=2;l<=n;m<<=1,l<<=1)
    {
        yy wn=yy(cos(pi/m),ops*sin(pi/m));
        for(int i=0;i<n;i+=l)
        {
            yy w=yy(1,0);
            for(int k=0;k<m;k++,w=w*wn)
            {
                yy t=a[i+k+m]*w;
                a[i+k+m]=a[i+k]-t;
                a[i+k]=a[i+k]+t;
            }
        }
    }
    if(ops==-1)
    for(int i=0;i<n;i++)
    a[i].x/=n;
}

yy a[N],b[N],c[N];
yy aa[N];
LL ans[N];

int main()
{
    cin>>n;
    for(int i=1;i<=n;i++)
    {
        int hy;
        scanf("%d",&hy);
        ans[hy]++;
        a[hy].x+=1.0;
        b[hy+hy].x+=1.0;
        c[hy+hy+hy].x+=1.0;
    }
    init(120010);
    fft(a,1);
    for(int i=0;i<n;i++)
    aa[i]=a[i]*a[i];
    fft(aa,-1);

    for(int i=0;i<n;i++)
    ans[i]+=(LL)(aa[i].x-b[i].x+0.4)/2;

    for(int i=0;i<n;i++)
    aa[i]=a[i]*a[i]*a[i];
    fft(aa,-1);

    fft(b,1);
    for(int i=0;i<n;i++)
    b[i]=b[i]*a[i];
    fft(b,-1);

    for(int i=0;i<n;i++)
    ans[i]+=(LL)(aa[i].x-3*b[i].x+2*c[i].x+0.4)/6;

    for(int i=0;i<n;i++)
    if(ans[i]!=0)
    printf("%d %lld\n",i,ans[i]);

    return 0;
}