[SP8372-TSUM]Triple Sums

题面在这里

description

某\(B\)姓\(OJ\)权限题

给出\(n\)个正整数\(a[i]\),求\(i<j<k\)且\(S=a[i]+a[j]+a[k]\)的三元组\((i,j,k)\)的个数

需要对所有可能的\(S\)进行求解

data range

\[n\le 40000,|a[i]|\le 20000
\]

solution

考虑构造生成函数求解

构造\(A(x)=\sum_{i=1}^{n}x^{a[i]},B(x)=\sum_{i=1}^{n}x^{2\times a[i]},C(x)=\sum_{i=1}^{n}x^{3\times a[i]}\)

然后容斥:$$f(x)=\frac{A(x)^3-3A(x)B(x)+2C(x)}{6}$$

code

这里本人写得复杂

#include<bits/stdc++.h>
#include<algorithm>
#include<iostream>
#include<cstdlib>
#include<iomanip>
#include<cstring>
#include<complex>
#include<vector>
#include<cstdio>
#include<string>
#include<bitset>
#include<ctime>
#include<cmath>
#include<queue>
#include<stack>
#include<map>
#include<set>
#define FILE "a"
#define mp make_pair
#define pb push_back
#define RG register
#define il inline
using namespace std;
typedef unsigned long long ull;
typedef vector<int>VI;
typedef long long ll;
typedef double dd;
const dd eps=1e-10;
const int mod=1e9+7;
const int N=1000010;
const dd pi=acos(-1);
const int inf=2147483647;
const ll INF=1e18+1;
il ll read(){
	RG ll data=0,w=1;RG char ch=getchar();
	while(ch!='-'&&(ch<'0'||ch>'9'))ch=getchar();
	if(ch=='-')w=-1,ch=getchar();
	while(ch<='9'&&ch>='0')data=data*10+ch-48,ch=getchar();
	return data*w;
}

il void file(){
	srand(time(NULL)+rand());
	freopen(FILE".in","r",stdin);
	freopen(FILE".out","w",stdout);
}

struct point{dd r,i;}a[N],b[N];
il point operator +(point a,point b){return (point){a.r+b.r,a.i+b.i};}
il point operator -(point a,point b){return (point){a.r-b.r,a.i-b.i};}
il point operator *(point a,point b){
	return (point){a.r*b.r-a.i*b.i,a.r*b.i+a.i*b.r};
}

int r[N],l;
il void FFT(point *a,int n,int opt){
	for(RG int i=0;i<n;i++)if(i<r[i])swap(a[i],a[r[i]]);
	for(RG int i=2;i<=n;i<<=1){
		point wn=(point){cos(2*pi/i),opt*sin(2*pi/i)};
		for(RG int j=0;j<n;j+=i){
			point w=(point){1,0};
			for(RG int k=j;k<j+(i>>1);k++,w=w*wn){
				point x=a[k+(i>>1)]*w;
				a[k+(i>>1)]=a[k]-x;
				a[k]=a[k]+x;
			}
		}
	}
	if(opt==-1)for(RG int i=0;i<n;i++)a[i].r/=n;
}

int n,lim,m,len,A[N];ll f[N],g[N],h[N];
int main()
{
	n=read();lim=inf;m=120000;for(len=1,l=0;len<=m;len<<=1,l++);
	for(RG int i=0;i<len;i++)r[i]=(r[i>>1]>>1)|((i&1)<<(l-1));
	for(RG int i=1;i<=n;i++)A[i]=read();lim=-20000;
	for(RG int i=1;i<=n;i++){
		f[A[i]-lim]++;g[2*A[i]-2*lim]++;h[3*A[i]-3*lim]++;
	}
	for(RG int i=0;i<len;i++)a[i].r=g[i],b[i].r=f[i];
	FFT(a,len,1);FFT(b,len,1);
	for(RG int i=0;i<len;i++)a[i]=a[i]*b[i];
	FFT(a,len,-1);
	for(RG int i=0;i<len;i++)g[i]=(ll)(a[i].r+0.5);

	for(RG int i=0;i<len;i++)a[i].r=f[i],a[i].i=0;
	FFT(a,len,1);
	for(RG int i=0;i<len;i++)a[i]=a[i]*a[i]*a[i];
	FFT(a,len,-1);
	for(RG int i=0;i<len;i++)f[i]=(ll)(a[i].r+0.5);

	for(RG int i=0;i<=m;i++)
		if((f[i]-3*g[i]+2*h[i])/6>0)
			printf("%d : %lld\n",i+3*lim,(f[i]-3*g[i]+2*h[i])/6);
	return 0;
}