ABC256E Kth Number题解

题意：给定 $p \leq n$ 个值域为 $1 \sim m$ 的数，你需要将其补充到 $n$ 个，每个数都在 $[1, m]$ 内独立随机。求最终第 $K$ 小数的期望值。 $n, m \leq 2000$

看到这个问题，就容易想到枚举第 $K$ 位的值，计算概率（或者方案数）。首先将其给定的 $p$ 个数用 $c n t$ 信息进行统计。假设第 $K$ 位选 $x$ ， $< x$ 的数有 $a$ 个， $= x$ 的数有 $b$ 个，则方案数为：

(\binom{p}{a - c n t_{< x}}) (\binom{p - (a - c n t_{< x})}{b - c n t_{x}}) \cdot (x - 1)^{a - c n t_{< x}} \cdot (m - x)^{n - a - b - c n t_{> x}}

这样统计是 $O (n^{3})$ 的，无法通过，且难以优化。我被这个思路卡了非常长时间，陷入了巨坑。

事实上，可以换一种思路统计，计算第 $K$ 位 $\geq x$ 的概率 $P (x)$ ，则等于 $x$ 的概率就是 $P (x) - P (x + 1)$ 。计算概率也不难，枚举有多少个 $< x$ ，假设为 $a$ 个，则此时的方案数为：

(\binom{p}{a - c n t_{< x}}) \cdot (x - 1)^{a - c n t_{< x}} \cdot (m - x + 1)^{n - a - c n t_{\geq x}}

除以总方案数即可得到 $P (x)$ 。注意枚举的时候不仅要保证 $a < k$ ，还要保证 $c n t_{\geq x} \leq n - a$ ，非常容易漏掉。

By cxm1024

#include <bits/stdc++.h>
using namespace std;
#define int long long
const int mod = 998244353;
const int MAXV = 2000;
int inv[MAXV + 10], jc[MAXV + 10], invjc[MAXV + 10];
int ksm(int a, int b, int res = 1) {
	for (; b; a = a * a % mod, b >>= 1)
		if (b & 1) res = res * a % mod;
	return res;
}
void init() {
	jc[0] = 1;
	for (int i = 1; i <= MAXV; i++)
		jc[i] = jc[i - 1] * i % mod;
	invjc[MAXV] = ksm(jc[MAXV], mod - 2);
	for (int i = MAXV; i > 0; i--)
		invjc[i - 1] = invjc[i] * i % mod;
	for (int i = 1; i <= MAXV; i++)
		inv[i] = jc[i - 1] * invjc[i] % mod;
}
int C(int x, int y) {
	return jc[x] * invjc[y] % mod * invjc[x - y] % mod;
}
int a[2010], p[2010], cnt[2010], mi[2010][2010];
signed main() {
	init();
	int n, m, k;
	cin >> n >> m >> k;
	for (int i = 1; i <= n; i++) {
		cin >> a[i];
		if (a[i] != 0) cnt[a[i]]++;
	}
	for (int i = 1; i <= m; i++)
		cnt[i] += cnt[i - 1];
	int P = ksm(ksm(m, n - cnt[m]), mod - 2);
	for (int i = 0; i <= 2000; i++) {
		mi[i][0] = 1;
		for (int j = 1; j <= 2000; j++)
			mi[i][j] = mi[i][j - 1] * i % mod;
	}
	for (int i = 1; i <= m; i++) {
		int res = 0;
		for (int j = cnt[i - 1]; j < k; j++)
			if (cnt[m] - cnt[i - 1] <= n - j)
				(res += mi[i - 1][j - cnt[i - 1]] * mi[m - i + 1][n - j - (cnt[m] - cnt[i - 1])] % mod * C(n - cnt[m], j - cnt[i - 1]) % mod) %= mod;
		p[i] = res * P % mod;
	}
	int ans = 0;
	for (int i = 1; i <= m; i++)
		(ans += i * (p[i] - p[i + 1] + mod) % mod) %= mod;
	cout << ans << endl;
	return 0;
}