Pythonで式の結果の最大出現頻度の期待値を求めるプログラム

問題の概要

M個の異なる式があり、それぞれの答えは1からNまでの範囲（両端を含む）に収まるとします。ここで、1からNまでの各iについての出現頻度を f(i) としたとき、x = max(f(i)) の期待値 E(x) を求めるのがこの問題です。

具体例で理解する

たとえば入力が M = 3、N = 3 の場合、答えは 2.2 になります。3つの式が取りうるすべての出力列と、それぞれの場合の最大出現頻度を表にまとめると、次のようになります。

全10通りのうち、最大頻度が1になるのは1通り、2になるのは6通り、3になるのは3通りです。したがって、期待値は次のように計算できます。

$$E(x) = \sum P(x) \times x = P(1) + 2P(2) + 3P(3) = \frac{1}{10} + 2 \times \frac{6}{10} + 3 \times \frac{3}{10} = \frac{22}{10}$$

解法のアプローチ

この問題は、二項係数（組み合わせの数 nCr）の計算結果をメモ化（キャッシュ）しながら、包除原理に基づく集計を行うことで効率的に解けます。手順は以下のとおりです。

• combination := 計算済みの二項係数を保存するための新しいマップ（辞書）
• 関数 nCr() を定義する。引数は n と k_in
• k := min(k_in, n − k_in)
• n < k または k < 0 の場合は 0 を返す
• (n, k) がすでに combination に存在する場合は、キャッシュされた値 combination[n, k] を返す
• k == 0 の場合は 1 を返す
• n == k の場合も 1 を返す
• それ以外の場合：
  • a := 1
  • cnt を 0 から k − 1 まで繰り返す：
    • a := a × (n − cnt)
    • a := a ÷ (cnt + 1) の切り捨て
    • combination[n, cnt + 1] := a
  • a を返す
• メイン処理では以下を実行する：
• arr := 新しいリスト
• k を 2 から M + 1 まで繰り返す：
  • a := 1、s := 0
  • i を 0 から ⌊M / k⌋ + 2 まで繰り返す：
    • M < i × k ならばループを抜ける
    • s := s + a × nCr(N, i) × nCr(N − 1 + M − i×k, M − i×k)
    • a := −a
  • arr の末尾に s を追加する
• total := arr の最後の要素
• diff := 先頭に arr[0] を置き、続けて (arr[cnt + 1] − arr[cnt])（cnt は 0 から M − 2）を並べた配列
• output := Σ diff[cnt] × (cnt + 1) / total（cnt は 0 から M − 1）
• output を返す

実装例

理解を深めるために、以下のPython実装を見てみましょう。

combination = {}
def nCr(n, k_in):
    k = min(k_in, n - k_in)
    if n < k or k < 0:
        return 0
    elif (n, k) in combination:
        return combination[(n, k)]
    elif k == 0:
        return 1
    elif n == k:
        return 1
    else:
        a = 1
        for cnt in range(k):
            a *= (n - cnt)
            a //= (cnt + 1)
            combination[(n, cnt + 1)] = a
        return a

def solve(M, N):
    arr = []
    for k in range(2, M + 2):
        a = 1
        s = 0
        for i in range(M // k + 2):
            if (M < i * k):
                break
            s += a * nCr(N, i) * nCr(N - 1 + M - i * k, M - i * k)
            a *= -1
        arr.append(s)
    total = arr[-1]
    diff = [arr[0]] + [arr[cnt + 1] - arr[cnt] for cnt in range(M - 1)]
    output = sum(diff[cnt] * (cnt + 1) / total for cnt in range(M))
    return output

M = 3
N = 3
print(solve(M, N))

入力

M = 3, N = 3

出力

2.2

まとめ

このアルゴリズムでは、二項係数 nCr を辞書にキャッシュすることで同じ計算の繰り返しを避け、包除原理（正負の項を交互に足し合わせる手法）によって各最大頻度に対応する場合の数を数えています。これにより、全列挙よりも大幅に少ない計算量で期待値を求めることが可能になります。