ABC256E Kth Number题解

题目传送门

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

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

\[\binom{p}{a-cnt_{<x}}\binom{p-(a-cnt_{<x})}{b-cnt_x}\cdot(x-1)^{a-cnt_{<x}}\cdot(m-x)^{n-a-b-cnt_{>x}} \]

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

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

\[\binom{p}{a-cnt_{<x}}\cdot(x-1)^{a-cnt_{<x}}\cdot(m-x+1)^{n-a-cnt_{\ge x}} \]

除以总方案数即可得到 \(P(x)\)。注意枚举的时候不仅要保证 \(a<k\)，还要保证 \(cnt_{\ge x}\le 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;
}

官方题解给出了另一种思考方向：

某个值的期望不仅等于 \(\sum\limits_{x}x\cdot P(x)\)，还等于 \(\sum\limits_{x}P(\ge x)\)（证明是容易的）。有了这个式子后，自然转化为了与上面相同的问题。