【luogu P5339】唱、跳、rap和篮球（容斥）（NTT）（OGF）

唱、跳、rap和篮球

题目链接：luogu P5339

题目大意

给你 a,b,c,d 四个字符的个数，然后问你可以组成多少个长度为 n 的字符串使得没有 abcd 这个子串。

思路

考虑容斥，用至少 \(0\) 个 \(abcd\) 子串减去至少 \(1\) 个 \(abcd\) 子串加上两个的……

然后考虑怎么排这些子串，可以发现这些子串不可能重叠（毕竟都是不同的），那我们可以把它看做一个整体 \(e\)。
那我们就变成了如果是至少 \(i\) 个，这些摆的方案数就是 \(\binom{n-3i}{i}\)。
（\(-3i\) 就相当于你把原本长度为四的 \(abcd\) 压缩为 \(e\)，每个少了三个字符，然后你就随便选 \(i\) 个位置放 \(e\)）

然后你接着考虑剩下的位置怎么随便放。
那不难想出它可以用一般生成函数 OGF 的思想来搞。
那我们考虑乘起来的四个多项式要如何构造。
考虑用有相同元素的排列数 \(\dfrac{n!}{n_1!n_2!...}\)

那我们四个多项式可以构造为形如 \(\dfrac{1}{a!},\dfrac{1}{a!},...,\dfrac{1}{a!},0,0,...\) 的形式。
（其中 \(\dfrac{1}{a!}\) 的最高项位置就是它能放的 \(a\) 个数，\(bcd\) 同理）
然后前面乘上一个 \((n-3i)!\) 就是每个的答案了。

代码

#include<cstdio>
#include<iostream>
#include<algorithm>
#define mo 998244353

using namespace std;

int n, a, b, c, d, ans, di, m;
int jc[20001], inv[20001], G, Gv;
int limit, l_size, an[20001];
int A[20001], B[20001], C[20001], D[20001];

int jia(int x, int y) {
	return (x + y >= mo) ? x + y - mo : x + y;
}

int jian(int x, int y) {
	return (x < y) ? (x - y + mo) : (x - y);
}

int cheng(int x, int y) {
	return 1ll * x * y % mo;
}

int CC(int n, int m) {
	if (m < 0 || m > n) return 0;
	return cheng(jc[n], cheng(inv[m], inv[n - m]));
}

int ksm(int x, int y) {
	int re = 1;
	while (y) {
		if (y & 1) re = cheng(re, x);
		x = cheng(x, x); y >>= 1;
	}
	return re;
}

void NTT(int *f, int op) {
	for (int i = 0; i < limit; i++)
		if (i < an[i]) swap(f[i], f[an[i]]);
	for (int mid = 1; mid < limit; mid <<= 1) {
		int Wn = ksm((op == 1) ? G : Gv, (mo - 1) / (mid << 1));
		for (int R = (mid << 1), j = 0; j < limit; j += R) {
			int w = 1;
			for (int k = 0; k < mid; k++, w = cheng(w, Wn)) {
				int x = f[j + k], y = cheng(w, f[j + mid + k]);
				f[j + k] = jia(x, y); f[j + mid + k] = jian(x, y);
			}
		}
	}
	if (op == -1) {
		int invl = ksm(limit, mo - 2);
		for (int i = 0; i < limit; i++) f[i] = cheng(f[i], invl);
	}
}

int P(int n, int a, int b, int c, int d) {
	if (n < 0 || a + b + c + d < n) return 0;
	limit = 1; l_size = 0;
	while (limit <= (a + b + c + d) * 2) {
		limit <<= 1; l_size++;
	}
	for (int i = 0; i < limit; i++)
		an[i] = (an[i >> 1] >> 1) | ((i & 1) << (l_size - 1));
	
	for (int i = 0; i < limit; i++) {
		A[i] = (i <= a) ? inv[i] : 0;
		B[i] = (i <= b) ? inv[i] : 0;
		C[i] = (i <= c) ? inv[i] : 0;
		D[i] = (i <= d) ? inv[i] : 0; 
	}
	NTT(A, 1); NTT(B, 1); NTT(C, 1); NTT(D, 1);
	for (int i = 0; i < limit; i++)
		A[i] = cheng(cheng(A[i], B[i]), cheng(C[i], D[i]));
	NTT(A, -1);
	
	return cheng(jc[n], A[n]);
}

int main() {
	G = 3; Gv = ksm(G, mo - 2);
	jc[0] = 1; for (int i = 1; i <= 20000; i++) jc[i] = cheng(jc[i - 1], i);
	inv[0] = inv[1] = 1; for (int i = 2; i <= 20000; i++) inv[i] = cheng(inv[mo % i], mo - mo / i);
	for (int i = 1; i <= 20000; i++) inv[i] = cheng(inv[i - 1], inv[i]);
	
	scanf("%d %d %d %d %d", &n, &a, &b, &c, &d);
	m = min(min(a, b), min(c, d));
	
	ans = 0; di = 1;
	for (int i = 0; i <= m; i++) {
		ans = jia(ans, cheng(cheng(di, CC(n - 3 * i, i)), P(n - 4 * i, a - i, b - i, c - i, d - i)));
		di = mo - di;
	}
	printf("%d", ans);
	
	return 0;
}