ABC295 E問題 Kth Number

期待値の求め方

昇順に並び替えたときの $K$ 番目の値が $x$ である確率を $P(x)$ とすると、
求めたい期待値は、
$\sum_{x=1}^{M} x P(x)$
です。

与えられた数列の中で $A_i = 0$ を満たす $i$ の個数を $p$ 個、 $0 < A_i < x$ の個数を $q$ 個、 $A_i = x$ の個数を $r$ 個、 $A_i > x$ の個数を $s$ 個とします( $p+q+r+s=N$ )。
$p$ 個の $A_i = 0$ を1以上 $M$ 以下の整数に置き換える場合の数は、 $M^p$ です。
このうち、 $K$ 番目の値が $x$ になるような置き換え方を考えます。 $A_i < x$ に置き換えられる個数を $t$ 個、 $A_i = x$ の個数を $u$ 個、 $A_i > x$ の個数を $v$ 個とします( $t+u+v=p$ )。このとき、昇順に並び替えると値が $x$ なのは、 $q + t + 1$ 番目から $q + r + t + u$ 番目までです。 $K$ がこの範囲に含まれればよいので、
$q + t + 1 \le K \le q + r + t + u$
これを満たす $t,u$ の範囲を考えればよいです。ある $t,u$ について考えると、 $p$ 個の中から $t$ 個と $u$ 個を選ぶ場合の数は、
${}_p \mathrm{C}_t \cdot {}_{p-t} \mathrm{C}_u = \frac{p!}{t!(p-t)!}\frac{(p-t)!}{u!(p-t-u)!}=\frac{p!}{t!u!(p-t-u)!}$
で、 $A_i < x$ への置き換えは1～ $x-1$ までの $x-1$ 通り、 $A_i > x$ への置き換えは $x+1$ ～ $M$ までの $M-x$ 通りあるので、場合の数は、
$\frac{p!}{t!u!(p-t-u)!}\left(x-1\right)^t\left(M-x\right)^{p-t-u}$
となります。これを条件を満たす $t,u$ の範囲で和を求めることで、昇順に並び替えたとき $K$ 番目の値が $x$ になる場合の数が求められ、これを $M^p$ で割ることで確率 $P(x)$ が求められます。

この考え方では、 $t,u$ のループで $P(x)$ を求め、 $x$ のループで答えを求める三重ループになりますが、実際には計算時間がかかりすぎます。

ループを減らす

これまで、 $K$ 番目の値がちょうど $x$ になる確率 $P(x)$ を考えていました。この発想を変えて、 $K$ 番目の値が $x$ 以上になる確率 $S(x)$ を考えてみます。
$S(x)=\sum_{k=x}^{M} P(k)$
この式から、
$P(x) = S(x) - S(x+1)$
となります。求めたい期待値は、 $S(x)$ で表すと、
$\sum_{x=1}^{M} x P(x)　= \sum_{x=1}^{M} x (S(x) - S(x+1))= \sum_{x=1}^{M} x S(x) -\sum_{x=2}^{M+1} (x-1) S(x)= \sum_{x=1}^{M} S(x)$
になります。
与えられた数列の中で $A_i = 0$ を満たす $i$ の個数を $p$ 個、 $0 < A_i < x$ の個数を $q$ 個、 $A_i \ge x$ の個数を $r$ 個とします( $p+q+r=N$ )。
$p$ 個の $A_i = 0$ を1以上 $M$ 以下の整数に置き換える場合の数は、 $M^p$ です。
このうち、 $K$ 番目の値が $x$ 以上になるような置き換え方を考えます。 $A_i < x$ に置き換えられる個数を $t$ 個、 $A_i \ge x$ の個数を $u$ 個とします( $t+u=p$ )。このとき、昇順に並び替えると値が $x$ 以上になるためには、 $q + t + 1 \le K$ であればよく、一重ループで済みます。
場合の数は、
$\frac{p!}{t!(p-t)!}\left(x-1\right)^t\left(M-x+1\right)^{p-t}$
となります。あとは、
$S(x) = \sum_{t=1}^{K-q-1} \frac{p!}{t!(p-t)!}\left(x-1\right)^t\left(M-x+1\right)^{p-t}$
から $\sum_{x=1}^{M} S(x)$ を求めることになります。

コード例 (Julia)

_MOD = 998244353

F = []
invF = []

function init_fact(N)
  push!(F,1)
  for i = 1:N
    push!(F,(F[i]*i)%_MOD)
  end
  for i = 1:N+1
    push!(invF,invmod(F[i],_MOD))
  end
end

function combination(N,A)
  if N < 0 || A > N || A < 0
    return 0
  end
  x = (F[N+1] * invF[A+1]) % _MOD
  x = (x * invF[N-A+1]) % _MOD
  return x
end

function solve()

  # 入力
  readInts() =  parse.(Int,split(readline()))
  readInt()  =  parse(Int,readline())

  N,M,K = readInts()
  A = readInts()

  B = []
  nzero = 0
  for i = 1:N
    if A[i] > 0
      push!(B,A[i])
    else
      nzero += 1    # A_i = 0 の個数
    end
  end

  # 階乗をあらかじめ計算しておく
  init_fact(nzero)

  # 0 以外の数字を昇順で並び替え
  sort!(B)

  ans = 0
  for k = 1:M
    # K 番目の数字が k 以上になる確率

    # K 番目の数字が k以上 <==> k より小さい数が K-1 個以下

    kpos1 = searchsortedfirst(B,k)  # k <= B[kpos1]

    # 「k より小さい数が K-1 個以下」 のうち、値が決まっている B を除いた個数
    under_max = (K-1) - (kpos1-1)

    for under = 0:under_max

      C = combination(nzero, under) # nzero 個の中から under個を選ぶ選び方
      n = 1
      if k-1>=0
        n = powermod( k-1, under, _MOD)   # under個は、1～k-1 の中から自由に1つ選べる
      end
      m = 1
      if M-k+1>=0
        m = powermod( M-k+1, nzero-under, _MOD)   # 残りは、k～M の中から自由に1つ選べる
      end

      s = (n * m) % _MOD
      s = (s * C) % _MOD
      ans += s
      ans = ans % _MOD
    end
  end

  ans = ans * invmod(powermod(M,nzero,_MOD), _MOD)
  ans = ans % _MOD

  println(ans)

end  # function solve

# main
solve()