【LG4491】[HAOI2018]染色

【LG4491】[HAOI2018]染色

题面

洛谷

题解

颜色的数量不超过\(lim=min(m,\frac nS)\)

考虑容斥，计算恰好出现\(S\)次的颜色至少\(i\)种的方案数\(f[i]\)，钦定\(i\)种颜色至少放\(S\)种

有\(m\)种颜色，那么要乘上\(C_m^i\)。

然后这\(n\)个位置分为\(i+1\)个部分：被钦定的\(i\)种颜色，每个\(S\)个；剩下\(m-i\)种颜色，一共\(n-iS\)种颜色，可以看作可重的全排列数，那么就有\(\frac{n!}{(S!)^i(n-iS)!}\)，但是后面情况的每个部分有\(m-i\)种取法，所以还要乘上\((m-i)^{n-iS}\)

\(\therefore f[i]=C_m^i\frac{n!}{(S!)^i(n-iS)!}(m-i)^{n-iS}\)

那么答案是什么呢？

\(ans[i]:\)恰好出现\(S\)次的颜色恰好\(i\)种的方案数

用容斥搞一下：

\[ans[i]=\sum_{j=i}^{lim}(-1)^{j-i}C_j^if[j]\\ =\sum_{j=i}^{lim}(-1)^{j-i}\frac{j!}{i!(j-i)!}f[j]\\ \Leftrightarrow ans[i]*i!=\sum_{j=i}^{lim}\frac{(-1)^{j-i}}{(j-i)!}f[j]*j! \]

然后\(NTT\)就可以了。

代码

#include <iostream>
#include <cstdio>
#include <cstdlib>
#include <cstring> 
#include <cmath> 
#include <algorithm>
using namespace std; 
inline int gi() {
    register int data = 0, w = 1;
    register char ch = 0;
    while (!isdigit(ch) && ch != '-') ch = getchar(); 
    if (ch == '-') w = -1, ch = getchar();
    while (isdigit(ch)) data = 10 * data + ch - '0', ch = getchar();
    return w * data; 
}
const int Mod = 1004535809;
int fpow(int x, int y) {
	int res = 1; 
	while (y) {
		if (y & 1) res = 1ll * res * x % Mod; 
		x = 1ll * x * x % Mod; 
		y >>= 1; 
	} 
	return res; 
} 
const int G = 3, iG = fpow(G, Mod - 2); 
const int MAX_N = 1e7 + 5, MAX_M = 1e5 + 5; 
int fac[MAX_N]; 
int C(int n, int m) {
	if (m > n) return 0; 
	return 1ll * fac[n] * fpow(fac[n - m], Mod - 2) % Mod * fpow(fac[m], Mod - 2) % Mod;
} 
int N, M, S, Limit, mx; 
int W[MAX_N], A[MAX_N << 2], B[MAX_N << 2], rev[MAX_N << 2]; 
void NTT(int *p, int op) { 
	for (int i = 0; i < Limit; i++) if (rev[i] > i) swap(p[rev[i]], p[i]); 
	for (int i = 1; i < Limit; i <<= 1) {
		int rot = fpow(op == 1 ? G : iG, (Mod - 1) / (i << 1)); 
		for (int j = 0, pls = (i << 1); j < Limit; j += pls) {
			int w = 1; 
			for (int k = 0; k < i; k++, w = 1ll * w * rot % Mod) {
				int x = p[j + k], y = 1ll * w * p[i + j + k] % Mod; 
				p[j + k] = (x + y) % Mod, p[i + j + k] = (x - y + Mod) % Mod; 
			} 
		} 
	} 
} 
int main () { 
	N = gi(), M = gi(), S = gi(); 
	for (int i = 0; i <= M; i++) W[i] = gi(); 
	mx = max(N, M); 
	fac[0] = 1; for (int i = 1; i <= mx; i++) fac[i] = 1ll * fac[i - 1] * i % Mod; 
	int mn = min(M, N / S); 
	Limit = 1; int P = 0; 
	while (Limit < (mn + 1) << 1) Limit <<= 1, ++P; 
	for (int i = 0; i <= mn; i++)
		A[i] = 1ll * fac[i] * C(M, i) % Mod *
			fac[N] % Mod * fpow(M - i, N - i * S) % Mod *
			fpow(1ll * fpow(fac[S], i) * fac[N - i * S] % Mod, Mod - 2) % Mod; 
	for (int i = 0; i <= mn; i++) {
		B[i] = fpow(fac[mn - i], Mod - 2); 
		if ((mn - i) & 1) B[i] = Mod - B[i]; 
	} 
	for (int i = 0; i < Limit; i++) rev[i] = (rev[i >> 1] >> 1) | ((i & 1) << (P - 1)); 
	NTT(A, 1), NTT(B, 1); 
	for (int i = 0; i < Limit; i++) A[i] = 1ll * A[i] * B[i] % Mod; 
	NTT(A, 0); 
	int inv = fpow(Limit, Mod - 2); 
	for (int i = 0; i < Limit; i++) A[i] = 1ll * inv * A[i] % Mod;
	int ans = 0; 
	for (int i = 0; i <= mn; i++) ans = (ans + 1ll * W[i] * A[mn + i] % Mod * fpow(fac[i], Mod - 2) % Mod) % Mod;
	printf("%d\n", ans); 
	return 0; 
}