CodeForces 958F3 Lightsabers (hard) 启发式合并/分治 多项式 FFT

原文链接http://www.cnblogs.com/zhouzhendong/p/8835443.html

题目传送门 - CodeForces 958F3

题意

　　有$n$个球，球有$m$种颜色，分别编号为$1\cdots m$，现在让你从中拿$k$个球，问拿到的球的颜色所构成的可重集合有多少种不同的可能。

　　注意同种颜色球是等价的，但是两个颜色为$x$的球不等价于一个。

　　$1\leq n\leq 2\times 10^5,\ \ \ \ \ 1\leq m,k\leq n$。

题解

　　来自Helvetic Coding Contest 2018 online mirror.

　　比赛的时候太蠢了只yy了个分治$FFT$，只有25分钟不敢写（其实说不定来得及，赛后写启发式合并不到20分钟A了（不过看了组数据（某种颜色出现次数为$0$的特殊情况）））。

　　分治$FFT$不讲，常数大容易被卡掉。

　　更好的做法是启发式合并。

　　考虑颜色集合S的计算结果为$a_{0\cdots x}$，其中$a_i$表示取$i$个球得到的不同结果数。

　　当合并两个颜色集合的时候，新的结果为：

$$c_i=\sum_{j=0}^{i}a_jb_{i-j}$$

　　显然就是一个多项式卷积直接$FFT$即可。

　　初始情况就是对于每一个颜色，设$cnt_i$为颜色$i$的出现次数，那么该颜色下标范围为$0\cdots cnt_i$，值全部为$1$。

　　然后我们只需要启发式合并几下就可以了。

　　启发式合并用小根堆来维护$vector$，保存的是编号，关键字是$.size()$，用$vector$来存储计算结果防$MLE$。

　　注意再开始的时候处理掉某种颜色出现次数为$0$的情况，不然会挂。

　　时间复杂度$O(n\log^2 n)$。

代码

#include <bits/stdc++.h>
using namespace std;
typedef long long LL;
const int N=1<<18,mod=1009;
double PI=acos(-1.0);
int n,m,k,tot[N];
struct C{
	double r,i;
	C(){}
	C(double a,double b){r=a,i=b;}
	C operator + (C x){return C(r+x.r,i+x.i);}
	C operator - (C x){return C(r-x.r,i-x.i);}
	C operator * (C x){return C(r*x.r-i*x.i,r*x.i+i*x.r);}
}w[N],A[N],B[N];
int R[N];
vector <int> colors[N<<1];
struct cmp{
	bool operator ()(int a,int b){
		return colors[a].size()>colors[b].size();
	}
};
priority_queue <int,vector<int>,cmp> heap;
void FFT(C a[],int n){
	for (int i=0;i<n;i++)
		if (i<R[i])
			swap(a[i],a[R[i]]);
	for (int t=n>>1,d=1;d<n;d<<=1,t>>=1)
		for (int i=0;i<n;i+=(d<<1))
			for (int j=0;j<d;j++){
				C tmp=w[t*j]*a[i+j+d];
				a[i+j+d]=a[i+j]-tmp;
				a[i+j]=a[i+j]+tmp;
			}
}
void FFT_times(vector <int> &a,vector <int> &b,vector <int> &c){
	int n,d;
	for (int i=0;i<a.size();i++)
		A[i]=C(a[i],0);
	for (int i=0;i<b.size();i++)
		B[i]=C(b[i],0);
	for (n=1,d=0;n<a.size()+b.size()-1;n<<=1,d++);
	for (int i=0;i<n;i++){
		R[i]=(R[i>>1]>>1)|((i&1)<<(d-1));
		w[i]=C(cos(2*PI*i/n),sin(2*PI*i/n));
	}
	for (int i=a.size();i<n;i++)
		A[i]=C(0,0);
	for (int i=b.size();i<n;i++)
		B[i]=C(0,0);
	FFT(A,n),FFT(B,n);
	for (int i=0;i<n;i++)
		A[i]=A[i]*B[i],w[i].i*=-1.0;
	FFT(A,n);
	c.clear();
	for (int i=0;i<=a.size()+b.size()-2;i++)
		c.push_back(((LL)(A[i].r/n+0.5))%mod);
}
int main(){
	scanf("%d%d%d",&n,&m,&k);
	for (int i=1,x;i<=n;i++)
		scanf("%d",&x),tot[x]++;
	while (!heap.empty())
		heap.pop();
	int size=0;
	for (int i=1;i<=m;i++){
		if (tot[i]==0)
			continue;
		colors[++size].clear();
		for (int j=0;j<=tot[i];j++)
			colors[size].push_back(1);
		heap.push(size);
	}
	while (heap.size()>=2){
		int x=heap.top();
		heap.pop();
		int y=heap.top();
		heap.pop();
		FFT_times(colors[x],colors[y],colors[++size]);
		colors[x].clear(),colors[y].clear();
		heap.push(size);
	}
	printf("%d",colors[size][k]);
	return 0;
}