AtCoder Beginner Contest 331 G - Collect Them All【概率期望+容斥+多项式】

题目链接：ABC331_G

写在前面

将来如果回顾这道题，建议自己看完题意一定先重新推一遍。如果还是不够熟练，多去做一些同类型的题目吧。

题意：

盒子里有 \(N\) 张卡片，每张卡片上写着一个数字，数字的范围是 \(1,...,M\)，写着数字 \(i\) 的卡片有 \(C_i\) 张\(（C_i>0）\)。有放回地抽取卡片，每次抽取一张，问抽到过全部 \(M\) 种数字的期望抽取次数是多少？
\((1\leq M\leq N\leq2\times10^5)\)

思路梳理：

概率期望

首先转化该题的概率期望部分。
对于这种“询问抽到过全部 \(M\) 种物品的期望次数”的期望类题目，经典的套路是将期望转化为概率的后缀和。

设 \(E(i)\) 为抽取 \(i\) 次恰好抽齐的期望次数， \(P(i)\) 为抽取 \(i\) 次恰好抽齐的概率，则有：

\(Ans = \sum\limits_{i=1}^{\infty}E(i) = \sum\limits_{i=1}^{\infty}P(i)\times i\)

设 \(suf(i)\) 为 \(P(i)\) 的后缀和，即 \(suf(i) = \sum\limits_{j=i}^{\infty}P(j)\)

于是又有：\(\sum\limits_{i=1}^{\infty}P(i)\times i = \sum\limits_{i=1}^{\infty}suf(i)\)

考虑 \(suf(i)\) 的实际含义，作为 \(P(i)\) 的后缀和，它表示抽取至少 \(i\) 次才抽齐的概率，也就是抽取 \(i-1\) 次都没有抽齐的概率。

设 \(P_1(i)\) 为抽取 \(i\) 次都没有抽齐的概率，于是得到：\(Ans = \sum\limits_{i=1}^{\infty}suf(i) = \sum\limits_{i=0}^{\infty}P_1(i)\)

容斥

现在来考虑 \(P_1(i)\) 该如何求得。先枚举抽取次数 \(i\) ，再枚举被抽到数字的集合 \(T\)（\(T\varsubsetneqq U\)，此处计算没有抽齐的情况，因此 \(T\) 为全集 \(U\) 的真子集）。设 \(f(i,T)\) 为抽取 \(i\) 次抽到数字的集合至多为 \(T\) （即没有抽到 \(T\) 集合外的数字）的概率， \(F(T)\) 为抽到数字的集合恰好为 \(T\) 的概率(包含抽取 \(0\sim \infty\) 次的情况），则答案为：

\(Ans = \sum\limits_{i=0}^{\infty}P_1(i) = \sum\limits_{i=0}^{\infty}\sum\limits_{T\varsubsetneqq U}^{\infty}f(i,T) = \sum\limits_{T\varsubsetneqq U}^{\infty}F(T)\)

现在将答案表示成了一个与集合有关的概率，但 \(F(T)\) 是一个不好计算的东西。“恰好”难于计算的情况下，可以考虑容斥的思想，将容易计算的“至多/至少”通过容斥的方法转化成难于计算的“恰好”。

设 \(G(T)\) 为抽到数字的集合至多为 \(T\) （即没有抽到 \(T\) 集合外的数字）的概率，发现 \(G(T)\) 是容易计算的:

\(G(T)=\sum\limits_{i=0}^{\infty}\left(\frac {1}{n} \times \sum\limits_{x\in T}C_x\right)^i\)

可以发现 \(G(T)\) 为等比级数求和，于是有：

\(G(T)= \frac {n}{n-\sum\limits_{x\in T}C_x}\)

建立“至多/至少”和“恰好”的关系，设 \(C(T)\) 为 \(F(T)\) 需要产生的贡献， \(coef(T)\) 为容斥系数，有容斥方法：

\(\sum\limits_{T\varsubsetneqq U}^{\infty}F(T)\times C(T) = \sum\limits_{T\varsubsetneqq U}^{\infty}G(T)\times coef(T)\)

由上面的答案计算式，可以得到 \(C(T) = 1 (C(U)=0)\)。因为“至多”包含“恰好”，\(G(T)\) 包含 \(F(T)\)， 在容斥中从包含的一方向被包含转移，得到 \(coef(T) = C(T) - \sum\limits_{S\varsubsetneqq T}coef(S)\)。

现在答案式为：

\(Ans = \sum\limits_{T\varsubsetneqq U}^{\infty}G(T)\times coef(T)\)

其中 \(G(T)\) 和 \(coef(T)\) 都可在有限复杂度内求出，消去了 \(\infty\)。

子集反演

根据子集/超集反演的知识：

若有

\(F(T) = \sum\limits_{S\subseteq T}G(S)\)

则有

\(G(S) = \sum\limits_{T\subseteq S}F(T)\times (-1)^{|T|-|S|}\)
\(G(T) = \sum\limits_{T\subseteq S}F(S)\times (-1)^{|T|-|S|}\)

在本题中，由\(coef(T) = C(T) - \sum\limits_{S\varsubsetneqq T}coef(S)\)，也就是 \(C(T) = \sum\limits_{S\subseteq T}coef(S)\)可以得到：

\(coef(T) = \sum\limits_{T\subseteq S}C(S)\times (-1)^{|T|-|S|}\)

又因为 \(C(T) = 1 (C(U)=0)\)，推导可知 \(coef(T) = (-1)^{M-|T|+1}\)

此时

\(Ans = \sum\limits_{T\varsubsetneqq U}^{\infty}G(T)\times (-1)^{M-|T|+1} = (-1)^{M+1} \times \sum\limits_{T\varsubsetneqq U}^{\infty}G(T)\times (-1)^{|T|} = (-1)^{M+1} \times \sum\limits_{T\varsubsetneqq U}^{\infty}\frac {\sum\limits_{x\in T}C_x}{n-\sum\limits_{x\in T}C_x}\times (-1)^{|T|}\)

背包、多项式

考虑枚举 \(\sum\limits_{x\in T}C_x\) 的值，集合对该值的贡献（即上式中的 \((-1)^{|T|}\)）可以用背包统计。而背包的过程可以用多项式乘法进行优化。

写了两版代码，一版用了Atcoder自带的多项式，一版手写了NTT。

Atc自带多项式

#include
#include
using namespace std;
using mint = atcoder::modint998244353;
const int N = 2e5 + 10;
int n, m, c[N];
int main() {
	cin >> n >> m;
	queue >q;
	for (int i = 1; i <= m; i++) {
		cin >> c[i];
		vectorf(c[i] + 1);
		f[0] = 1, f[c[i]] = -1;
		q.push(f);
	}
	while (q.size() > 1) {
		vectora = q.front();
		q.pop();
		vectorb = q.front();
		q.pop();
		q.push(atcoder::convolution(a, b));
	}
	mint ans = 0;
	vectorf = q.front();
	for (int i = 0; i < n; i++) {
		ans += mint(n) / (n - i) * f[i];
	}
	if (!(m & 1))ans = -ans;
	cout << ans.val() << endl;
}

手写NTT

#include
#define ll long long
using namespace std;
const int N = 4e5 + 10, G = 3, mod = 998244353;
int n, m, c[N];
ll invn, invg, f[N], g[N], nn, pos[N];
ll pw(ll x, ll k) {
	ll num = 1;
	while (k) {
		if (k & 1)num = num * x % mod;
		x = x * x % mod;
		k >>= 1;
	}
	return num;
}
void ntt(ll *f, bool flag, int nn) {
	invn = pw(nn, mod - 2);
	for (int i = 0; i <= nn; i++) {
		pos[i] = (pos[i >> 1] >> 1) | ((i & 1) ? nn >> 1 : 0);
	}
	for (int i = 0; i <= nn; i++) {
		if (i < pos[i])swap(f[i], f[pos[i]]);
	}
	for (int i = 2; i <= nn; i <<= 1) {
		ll fir = pw(flag ? G : invg, (mod - 1) / i);
		for (int j = 0, p = (i >> 1); j + i - 1 <= nn; j += i) {
			ll bur = 1;
			for (int k = j; k < j + p; k++) {
				ll tt = bur * f[k + p] % mod;
				f[k + p] = f[k] - tt, f[k + p] += (f[k + p] < 0 ? mod : 0);
				f[k] += tt, f[k] -= (f[k] >= mod ? mod : 0);
				bur = bur * fir % mod;
			}
		}
	}
	if (!flag) {
		for (int i = 0; i <= nn; i++)f[i] = f[i] * invn % mod;
	}
}
void mul(ll *a, ll *b, int n, int m) {
	for (nn = 1; nn <= n + m; nn <<= 1);
	ntt(a, 1, nn), ntt(b, 1, nn);
	for (int i = 0; i <= nn; i++)a[i] = a[i] * b[i] % mod;
	ntt(a, 0, nn);
}
int main() {
	cin >> n >> m;
	invg = pw(G, mod - 2);
	queue >q;
	for (int i = 1; i <= m; i++) {
		cin >> c[i];
		vectorf(c[i] + 1);
		f[0] = 1, f[c[i]] = -1;
		q.push(f);
	}
	while (q.size() > 1) {
		vectora = q.front();
		q.pop();
		vectorb = q.front();
		q.pop();
		for (int i = 0; i < a.size(); i++)f[i] = a[i];
		for (int i = 0; i < b.size(); i++)g[i] = b[i];
		mul(f, g, a.size() - 1, b.size() - 1);
		vectorc;
		for (int i = 0; i <= nn; i++) {
			c.push_back(f[i]);
			f[i] = 0, g[i] = 0;
		}
		while (c[c.size() - 1] == 0)c.pop_back();
		q.push(c);
	}
	ll ans = 0;
	vectorf = q.front();
	for (int i = 0; i < n; i++) {
		ans = (ans + 1ll * n * pw(n - i, mod - 2) % mod * f[i] % mod) % mod;
	}
	if (!(m & 1))ans = mod - ans;
	cout << ans << endl;
	return 0;
}