Official

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: