[题解]CF958F3 Lightsabers (hard)

Posted on Edited on In

题意

给出\(m\)个集合, 第\(i\)个集合有\(c_i\)种颜色为\(i\)的小球, 保证\(\sum_{i = 1}^m c_i = n\). 现从这些集合中取出\(k\)个小球, 问这些小球构成的集合有多少种. 答案模\(1009\).

\(1 \leq m,k \leq n \leq 2\cdot 10^5\)

题解

普通型生成函数裸题.

设\(G(x) = \prod_{i = 1}^m(\sum_{j = 0}^{c_i}x^j)\)

那么\(G(x)[k]\)即为答案. 问题转换为\(m\)个多项式相乘. 使用分治法可在\(O(n\log^2n)\)时间内用FFT求解.

代码

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
//
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
#define endl '\n'
const int maxn = 4e6+10;
const int p = 1009;
int raw[maxn],cnt[maxn];
typedef double ld;
typedef vector<int> Poly;
const ld PI = acos(-1.0);
int rev[maxn];
struct Comp{
    double x, y;
    Comp (double x = 0, double y = 0):x(x),y(y){}
    Comp operator * (Comp B){
   		return Comp(x*B.x - y*B.y,x*B.y + y*B.x);
	}
	Comp operator + (Comp B){
	    return Comp(x + B.x,y + B.y);
	}
	Comp operator - (Comp B){
	    return Comp(x - B.x,y - B.y);
	}
};
Comp F1[maxn],F2[maxn];
void FFT(Comp * A,int siz,int type){
	int n = siz;
	int S = log2(n);
	for(int i = 0;i < n;i++)rev[i] = (rev[i>>1] >> 1) | ((i & 1) << (S - 1));
	for(int i = 0;i < n;i++)if(i < rev[i])swap(A[i],A[rev[i]]);
	for(int i = 1;i < n;i *= 2){
		Comp Wn(cos(PI/i),type*sin(PI/i));
		for(int j = 0;j < n;j += i*2){
			Comp W(1,0);
			for(int k = 0;k < i;k++){
				Comp facx = A[j+k],facy = W*A[j+k+i];
				A[j+k] = facx + facy;
				A[j+k+i] = facx - facy;
				W = W*Wn;
			}
		}
	}
    if(type == -1)for(int i = 0;i < n;i++)A[i].x = ((A[i].x/n + 0.5));
}
Poly mul(Poly A,Poly B){
	int n = A.size(),m = B.size();
	int siz = n + m - 1;
	Poly C(siz);
	if(siz < 64){
		for(int i = 0;i < n;i++){
			for(int j = 0;j < m;j++)C[i+j] = (C[i+j] + 1LL*A[i]*B[j])%p;
		}
		return C;
	}
	int fsiz = 1;
	while(fsiz <= siz)fsiz *= 2;
	for(int i = 0;i < fsiz;i++)F1[i] = F2[i] = 0;
	for(int i = 0;i < n;i++)F1[i] = A[i];
	for(int i = 0;i < m;i++)F2[i] = B[i];
	FFT(F1,fsiz,1);
	FFT(F2,fsiz,1);
	for(int i = 0;i < fsiz;i++)F1[i] = F1[i]*F2[i];
	FFT(F1,fsiz,-1);
	for(int i = 0;i < siz;i++){
		C[i] = ((ll)F1[i].x)%p;
	}
	return C;
}
Poly solve(int l,int r){
	if(l == r){
		return Poly(cnt[l]+1,1);
	}
	int m = (l+r)/2;
	return mul(solve(l,m),solve(m+1,r));
}
Poly ans;
signed main(){
	ios::sync_with_stdio(0);
	cin.tie(0);
	int n,m,k;
	cin >> n >> m >> k;
	for(int i = 1;i <= n;i++){
		int x;
		cin >> x;
		cnt[x]++;
	}
	Poly ans = solve(1,m);
	cout << ans[k] << endl;
	return 0;
}