[BZOJ3771] Triple

Description

给定一个数列，要求从中选出不超过 (3) 个数，问每种和对应多少种组合。

Solution

设 (f) 是原数列对应的生成函数，(g) 是双倍（即每个元素使用 2 次）对应的生成函数，(h) 是三倍

考虑三个数的情况，先做排列，先取 (f*f*f)，然后减去 ABB,BAB,BBA 式的方案 (3(f*g-h))，再减去 AAA 式的方案 (h)，故方案数对应的生成函数为

[frac 1 6 (f*f*f-3f*g+2h) ]

用 FFT 计算即可

一个数的情况显然为 (f)，两个数的情况为 ((f*f-g)/2)

#include<bits/stdc++.h>
#define N 262145
#define pi acos(-1)
using namespace std;

namespace po {
typedef complex<double> E;
int n,m,L;
int R[N];
E a[N],b[N];

void fft(E *a,int f){
	for(int i=0;i<n;i++)if(i<R[i])swap(a[i],a[R[i]]);
	for(int i=1;i<n;i<<=1){
		E wn(cos(pi/i),f*sin(pi/i));
		for(int p=i<<1,j=0;j<n;j+=p){
			E w(1,0);
			for(int k=0;k<i;k++,w*=wn){
				E x=a[j+k],y=w*a[j+k+i];
				a[j+k]=x+y;a[j+k+i]=x-y;
			}
		}
	}
}

void mul(int _n,int *aa,int *bb,int *cc){
    n=_n;
    memset(a,0,sizeof a);
    memset(b,0,sizeof b);
    memset(R,0,sizeof R);
    L=0;
    m=n;
	for(int i=0,x;i<=n;i++)a[i]=aa[i];
	for(int i=0,x;i<=m;i++)b[i]=bb[i];
	//m*=2;
	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,1);fft(b,1);
	for(int i=0;i<=n;i++)a[i]=a[i]*b[i];
	fft(a,-1);
	memset(cc,0,sizeof cc);
	for(int i=0;i<=m;i++) cc[i]=(int)(a[i].real()/n+0.5);
}
}

using po::mul;

int n,k,a[N],b[N],c[N],f[N],g[N],h[N];

signed main() {
    ios::sync_with_stdio(false);
    int m=120000;
    cin>>n;
    for(int i=1;i<=n;i++) {
        int t;
        cin>>t;
        f[t]++;
        g[2*t]++;
        h[3*t]++;
    }
    mul(m,f,f,a);
    for(int i=1;i<=m;i++) {
        c[i]+=a[i]-g[i];
        c[i]/=2;
    }
    mul(m,a,f,a);
    mul(m,f,g,b);
    for(int i=1;i<=m;i++) {
        c[i]+=(a[i]-3*b[i]+2*h[i])/6;
        c[i]+=f[i];
        if(c[i]) cout<<i<<" "<<c[i]<<endl;
    }
}