E - Kth Number Editorial by en_translator
In general, for a probability variable \(X\) that takes an integer value between \(1\) and \(M\),
\[(\text{expected value of }X)=\sum_{i=1}^{M}(i\times (\text{probability that }X = i))=\sum_{i=1}^{M}(\text{probability that }X \geq i).\]
We can apply this formula to this problem as follows; For each \(i=1,2,\dots ,M\), it is sufficient to find the probability that \(A_K \geq i\) after the two operations. This equals the probability that, after operation 1, there are at least \((N+1-K)\) integers \(j\) such that \(A_j \geq i\).” Thus, we can divide into cases depending on how many \(0\)s in \(A\) is replaced by integers \(i\) or greater; then it is boiled down to the probability of repeated trails. The complexity can be \(O(NM)\) or \(O(NM\log N)\) depending on the implementation.
Sample code (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: