Official
E - Kth Number Editorial by yuto1115
解説一般に、\(1\) 以上 \(M\) 以下の整数値を取る確率変数 \(X\) について、
\[(X\ \text{の期待値})=\sum_{i=1}^{M}(i\times (X = i\ \text{となる確率}))=\sum_{i=1}^{M}(X \geq i\ \text{となる確率})\]
が成り立ちます。
本問題に対してこの式変形を適用すると、各 \(i=1,2,\dots ,M\) に対して「操作 1, 2 を行ったあと \(A_K \geq i\) となる確率」を求めればよいことになります。これは「操作 1 を行ったあと、\(A_j \geq i\) を満たす \(j\) が \(N+1-K\) 個以上ある確率」と言い換えられるので、\(A\) に含まれる \(0\) のうち操作 1 で \(i\) 以上に置き換えられるものの数を全探索すれば、反復試行の確率を求める問題に帰着されます。計算量は実装方針によって \(O(NM)\) や \(O(NM\log N)\) 等になります。
実装例 (C++) :
#include<bits/stdc++.h>
#include<atcoder/modint>
using namespace std;
using namespace atcoder;
using mint = modint998244353;
mint binom[5010][5010];
int main() {
for (int i = 0; i < 5010; i++) {
binom[i][0] = 1;
for (int j = 1; j <= i; j++) {
binom[i][j] = binom[i - 1][j - 1] + binom[i - 1][j];
}
}
int n, m, k;
cin >> n >> m >> k;
vector<int> a(n);
for (int &i: a) cin >> i;
mint ans = 0;
for (int i = 1; i <= m; i++) {
int rem = n + 1 - k;
int zero = 0;
for (int j: a) {
if (j >= i) --rem;
if (j == 0) ++zero;
}
if (rem < 0 or rem > zero) {
ans += (rem < 0 ? 1 : 0);
continue;
}
mint p = mint(m + 1 - i) / m;
vector<mint> p_pow(zero + 1, 1);
vector<mint> q_pow(zero + 1, 1);
for (int j = 0; j < zero; j++) {
p_pow[j + 1] = p_pow[j] * p;
q_pow[j + 1] = q_pow[j] * (1 - p);
}
for (int j = rem; j <= zero; j++) {
ans += binom[zero][j] * p_pow[j] * q_pow[zero - j];
}
}
cout << ans.val() << endl;
}
posted:
last update: