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: