Kazun の競プロ記録

競技プログラミングに関する様々な話題を執筆します.

AtCoder Beginner Contest 295 E問題 Kth Number

AtCoder ABC ABC295 500 pts E問題

問題

atcoder.jp

提出解答

atcoder.jp

問題の概要

A_1, \dots, A_N は 0 以上 M 以下の整数である.

ここで, A_1, \dots, A_N にある 0 の項全てに対して, 以下の操作を独立に行う.

  • 1 以上 M 以下の整数から一様ランダムに選び, その項をその整数で置き換える.

この後, A_1, \dots, A_N を昇順に並べた時の第 K 項の期待値を求めよ.

制約

  • 1 \leq K \leq N \leq 2000
  • 1 \leq M \leq 2000
  • 0 \leq A_i \leq M

解法

最後に並び替えるので, 必要ならば A を並び替えて, A_1, \dots, A_r \geq 0, A_{r+1}, \dots, A_N=0 としてもよい.

A_1, \dots, A_N を昇順に並べたときの第 K 項を表わす確率変数を X とする.

このとき, X の期待値 E[X ] は

\displaystyle \begin{align*} E[X ] &=\sum_{a=1}^M a \operatorname{Prob}(X=a) \end{align*}

となる. ここで, 1 \leq a \leq M であるから, \displaystyle a=\sum_{b=1}^M [b \leq a] となる. なお, [\bullet ] は アイバーソンの記法 (アイバーソンの記法 - Wikipedia) を表わす.

よって,

\displaystyle \begin{align*} E[X] &=\sum_{a=1}^M a \operatorname{Prob}(X=a) \\ &=\sum_{a=1}^M \sum_{b=1}^M [b \leq a ] \operatorname{Prob}(X=a) \\ &=\sum_{b=1}^M \left(\sum_{a=1}^M [b \leq a ] \operatorname{Prob}(X=a) \right) \\ \end{align*}

となる. ここで,

\displaystyle \sum_{a=1}^M [b \leq a ] \operatorname{Prob}(X=a)=\operatorname{Prob}(X \geq b)

であるから,

\displaystyle E[X]=\sum_{b=1}^M \operatorname{Prob}(X \geq b)

となる.

よって, \operatorname{Prob}(X \geq b) を求めれば良い. これは A_1, \dots, A_N の中に b 以上が K 個ある確率になる.

A_1, \dots, A_r の中に b 以上が s 個あるとする. このとき, A_1, \dots, \dots, A_N の中に b 以上が K 個あるようにするためには, A_{r+1}, \dots, A_N の中に b 以上が u_b:=\max(N-K+1-s, 0) 以上あることが必要十分条件になる.

従って, A_{r+1}, \dots, A_N の中にある b 以上の個数で場合分けすることにより,

\displaystyle \operatorname{Prob}(X \geq b)=\sum_{t=u_b}^{N-r} \dbinom{N-r}{t} \left(\dfrac{M-b+1}{M} \right)^t \left(\dfrac{b-1}{M} \right)^{N-r-t}

よって,

\displaystyle E[X]=\sum_{b=1}^M \left(\sum_{t=u_b}^{N-r} \dbinom{N-r}{t} \left(\dfrac{M-b+1}{M} \right)^t \left(\dfrac{b-1}{M} \right)^{N-r-t} \right)

となる.

これは計算量 O(NM \log N) 時間や O(NM) 時間で計算できる.