SP8372 Triple Sums [生成函数+容斥]

description

传送门
给$n$个数，让你输出所有可能的由三个不同下标的数得到的和，以及构成该和的方案数。

solution

很容易想到$A=cnt_0*x^0+cnt_1*x^1+cnt_2*x^2...$卷三次，但要减掉存在至少两个相同下标的方案。
因此构造$B=cnt_0*x^0+cnt_1*x^2+cnt_2*x^4...$，因此$A*B$即为三个或两个下标都相同的方案$*$1。
思考怎么容斥的时候分类讨论算凑目标贡献。
首先，熟悉的套路：选择的三元组可以按照排列做贡献，最后除以$6$。
$A$中三个下标相同的算了一次，而恰两个下标相同的算了三次。
因此$A-3*B$多减两个三个下标相同的，构造 $C=cnt_0*x^0+cnt_1*x^2+cnt_2*x^6...$。
所以最终答案为: $\dfrac{A-3AB+2C}{6}$

code

戳我

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef double db;
const int N=1e6+5;
const db pi=acos(-1);
struct cop {
	db img,rl;
	friend inline cop operator*(cop u,cop v) {return (cop){u.img*v.rl+u.rl*v.img,u.rl*v.rl-u.img*v.img};}
	friend inline cop operator+(cop u,cop v) {return (cop){u.img+v.img,u.rl+v.rl};}
	friend inline cop operator-(cop u,cop v) {return (cop){u.img-v.img,u.rl-v.rl};}
}A[N],B[N],C[N],K[N];
int up,rev[N];

void gt_up(int len) {
	up=1;int l=0;
	while(up<=len)up<<=1,l++;
	for(int i=1;i<up;i++)rev[i]=(rev[i>>1]>>1)|((i&1)<<(l-1));
}

void FFT(cop *a,int op) {
	for(int i=1;i<up;i++)if(rev[i]>i)swap(a[i],a[rev[i]]);
	for(int mid=1;mid<up;mid<<=1) {
		cop w1=(cop){op*sin(pi/mid),cos(pi/mid)};
		for(int len=mid<<1,l=0;l<up;l+=len) {
			cop w=(cop){0,1};
			for(int i=0;i<mid;i++,w=w*w1) {
				int p=l+i,q=p+mid;
				cop x=a[p],y=a[q]*w;
				a[p]=x+y;a[q]=x-y;
			}
		}
	}
}

int b[N],mn=114514,mx=-mn;

int main() {
	int n;scanf("%d",&n);
	for(int i=1;i<=n;i++) {scanf("%d",&b[i]);mx=max(mx,b[i]);mn=min(mn,b[i]);}
	if(mn<0) {
		mn*=-1;mx+=mn;
		for(int i=1;i<=n;i++)b[i]+=mn;
	}
	else mn=0;
//	printf("!");
	gt_up(mx*3);
	for(int i=1;i<=n;i++) {A[b[i]].rl++;B[b[i]<<1].rl+=3;C[b[i]*3].rl+=2;}
	FFT(A,-1);FFT(B,-1);FFT(C,-1);
	for(int i=0;i<up;i++)A[i]=A[i]*A[i]*A[i]-A[i]*B[i]+C[i];
	FFT(A,1);
	for(int i=0;i<3*mx;i++) {
		ll val=(ll)round(A[i].rl/up);
		if(!val)continue;
		printf("%d : %lld\n",i-3*mn,val/6);
	}
	return 0;
}

加一道类似的题： 【BZOJ3771】Triple