BZOJ 3771: Triple 生成函数 + FFT
Description
我们讲一个悲伤的故事。
从前有一个贫穷的樵夫在河边砍柴。
这时候河里出现了一个水神，夺过了他的斧头，说：
“这把斧头，是不是你的？”
樵夫一看：“是啊是啊！”
水神把斧头扔在一边，又拿起一个东西问：
“这把斧头，是不是你的？”
樵夫看不清楚，但又怕真的是自己的斧头，只好又答：“是啊是啊！”
水神又把手上的东西扔在一边，拿起第三个东西问：
“这把斧头，是不是你的？”
樵夫还是看不清楚，但是他觉得再这样下去他就没法砍柴了。
于是他又一次答：“是啊是啊！真的是！”
水神看着他，哈哈大笑道：
“你看看你现在的样子，真是丑陋！”
之后就消失了。
樵夫觉得很坑爹，他今天不仅没有砍到柴，还丢了一把斧头给那个水神。
于是他准备回家换一把斧头。
回家之后他才发现真正坑爹的事情才刚开始。
水神拿着的的确是他的斧头。
但是不一定是他拿出去的那把，还有可能是水神不知道怎么偷偷从他家里拿走的。
换句话说，水神可能拿走了他的一把，两把或者三把斧头。
樵夫觉得今天真是倒霉透了，但不管怎么样日子还得过。
他想统计他的损失。
樵夫的每一把斧头都有一个价值，不同斧头的价值不同。总损失就是丢掉的斧头价值和。
他想对于每个可能的总损失，计算有几种可能的方案。
注意：如果水神拿走了两把斧头a和b，(a,b)和(b,a)视为一种方案。拿走三把斧头时，(a,b,c),(b,c,a),(c,a,b),(c,b,a),(b,a,c),(a,c,b)视为一种方案。
Input
第一行是整数N，表示有N把斧头。
接下来n行升序输入N个数字Ai，表示每把斧头的价值。
Output
若干行，按升序对于所有可能的总损失输出一行x y，x为损失值，y为方案数。
题解：在 $n$ 个互不相同的数中选 $1$ 到 $3$ 个数，其和为 $sum$ 的方案有多少种
构造生成函数 $A=\sum_{i=1}^{Max}a_{i}x^i$ 表示是否有 $i$ 这个数 (0/1)
$B=\sum_{i=1}^{Max}a_{i}x^i$ 表示仅选两个相同数的生成函数 (0/1)
$C=\sum_{i=1}^{Max}a_{i}x^i$ 表示仅选三个相同数的生成函树 (0/1)
只选一个 ： $A_{sum}$
只选两个： $\frac{(A^2-B)}{2}$
只选三个：$\frac{A^3-3\times(A*B-C)-C}{6}$
依次解释一下：
(1) 方案 = 选一个
(2) 总 - 选两个相同 = 选两个不同的有序方案，方案 = 选两个不同有序方案/2
(3) (总 - 选两个相同一个不同有序方案 - 选三个相同) / 6
#include<bits/stdc++.h>
#define setIO(s) freopen(s".in","r",stdin), freopen(s".out","w",stdout)
#define ll long long
#define maxn 500000
using namespace std;
const double pi=acos(-1.0);
struct cpx
{
double x,y;
cpx(double a=0,double b=0) {x=a,y=b; }
cpx operator+(const cpx b) { return cpx(x+b.x,y+b.y); }
cpx operator-(const cpx b) { return cpx(x-b.x,y-b.y); }
cpx operator*(const cpx b) { return cpx(x*b.x-y*b.y,x*b.y+y*b.x); }
}A[maxn],B[maxn],C[maxn],D[maxn];
void FFT(cpx *a,int n,int flag)
{
for(int i=0,k=0;i<n;++i)
{
if(i>k) swap(a[i],a[k]);
for(int j=(n>>1);(k^=j)<j;j>>=1);
}
for(int mid=1;mid<n;mid<<=1)
{
cpx wn(cos(pi/mid), flag*sin(pi/mid)), x,y;
for(int i=0;i<n;i+=(mid<<1))
{
cpx w(1,0);
for(int j=0;j<mid;++j)
{
x=a[i+j], y=w*a[i+j+mid];
a[i+j]=x+y, a[i+j+mid]=x-y;
w=w*wn;
}
}
}
if(flag==-1) for(int i=0;i<n;++i) a[i].x/=(double) n;
}
int Max;
int arr[maxn],f[maxn],g[maxn*2],h[maxn*3];
ll answer[maxn];
inline void solve1()
{
for(int i=1;i<=Max;++i) answer[i]+=f[i];
}
inline void solve2()
{
int len=1;
for(int i=1;i<=Max;++i) A[i].x=f[i];
for(;len<=(Max<<1);len<<=1);
FFT(A,len,1);
for(int i=0;i<len;++i) A[i]=A[i]*A[i];
FFT(A,len,-1);
for(int i=0;i<len;++i) answer[i]+=((ll)(A[i].x+0.5)-g[i])/2;
for(int i=0;i<len;++i) A[i].x=A[i].y=0;
}
inline void solve3()
{
int len=1;
for(;len<=(Max*3);len<<=1);
for(int i=0;i<=Max;++i) A[i].x=f[i];
for(int i=0;i<=Max*3;++i) C[i].x=A[i].x, D[i].x=g[i];
FFT(A, len, 1);
for(int i=0;i<len;++i) A[i]=A[i]*A[i]*A[i];
FFT(A, len, -1);
FFT(C, len, 1), FFT(D, len, 1);
for(int i=0;i<len;++i) C[i]=C[i]*D[i];
FFT(C, len, -1);
for(int i=0;i<len;++i) answer[i]+=(ll)((ll)(A[i].x+0.5)-3*((ll)(C[i].x+0.5)-1ll*h[i])-1ll*h[i])/6;
}
int main()
{
// setIO("input");
int n;
scanf("%d",&n);
for(int i=1;i<=n;++i)
{
scanf("%d",&arr[i]);
++f[arr[i]], ++g[arr[i]<<1], ++h[arr[i]*3];
Max=max(Max, arr[i]);
}
solve1(),solve2(), solve3();
for(int i=0;i<3*Max;++i) if(answer[i]) printf("%d %lld\n",i, answer[i]);
return 0;
}
