Pythonで最大の平均合格率を求めるプログラム（ヒープを使った貪欲法）

問題の概要

クラスのリストが与えられ、classes[i] は [pass_i, total_i] という形式で表されます。ここで pass_i は i 番目のクラスで試験に合格した生徒の数、total_i はそのクラスの生徒の総数です。さらに extra という値も与えられます。これは、どのクラスに配属されても必ず試験に合格できることが保証されている「優秀な追加生徒」の人数を意味します。

私たちの課題は、これらの追加生徒を各クラスへ割り当てることで、全クラスの平均合格率を最大化することです。クラスの合格率は「そのクラスで合格する生徒数 ÷ クラスの総生徒数」で求められ、平均合格率は「全クラスの合格率の合計 ÷ クラス数」で計算されます。追加生徒をすべて割り当てた後の、最大可能な平均合格率を求めてください。

入力例

classes = [[2,3],[4,6],[3,3]]、extra = 3 の場合、出力は 0.83809 になります。最初のクラスに2人、2番目のクラスに1人の追加生徒を配属すると比率が最大化され、このときの平均は (4/5 + 5/7 + 3/3) / 3 = 0.83809 となります。

解決アプローチ：貪欲法＋ヒープ

この問題は、貪欲法とヒープ（優先度付きキュー）を組み合わせると効率的に解けます。鍵となるのは「あるクラスに生徒を1人追加したとき、合格率がどれだけ改善されるか」という限界改善幅です。改善幅が最も大きいクラスから順に生徒を割り当てていくことで、全体の平均合格率を最大化できます。

Python の heapq は最小ヒープしかサポートしないため、改善幅にマイナスを掛けた値を優先度として格納します。こうすることで、ヒープの先頭には常に「改善幅が最大のクラス」が来るようになります。

具体的な手順は次のとおりです。

• 各クラス (a, b) について、生徒を1人追加した場合の合格率の変化量 a/b − (a+1)/(b+1) を優先度とするタプル (変化量, a, b) のリスト h を作成します
• h をヒープ化します
• extra が 0 になるまで、次の処理を繰り返します
  • ヒープの先頭要素 (v, a, b) を取り出します（＝現時点で改善幅が最大のクラス）
  • (a, b) を (a+1, b+1) に更新して、生徒を1人追加します
  • 更新後の変化量 −((a+1)/(b+1) − a/b) を計算し、新しいタプルとしてヒープに戻します
  • extra を1減らします
• 最後に、ヒープ内の全タプルから合格率の平均を計算して返します

Python実装例

import heapq

def solve(classes, extra):
    # 生徒を1人追加した際の合格率の変化量（マイナス値）を優先度とするタプルを作成
    h = [(a / b - (a + 1) / (b + 1), a, b) for a, b in classes]
    heapq.heapify(h)
    while extra:
        v, a, b = heapq.heappop(h)          # 改善幅が最大のクラスを取り出す
        a, b = a + 1, b + 1                 # 生徒を1人追加
        heapq.heappush(h, (-(a + 1) / (b + 1) + a / b, a, b))  # 新しい変化量で戻す
        extra -= 1
    return sum(a / b for v, a, b in h) / len(h)

classes = [[2,3],[4,6],[3,3]]
extra = 3
print(solve(classes, extra))

入力

[[2,3],[4,6],[3,3]], 3

出力

0.8380952380952381

結果は約 0.83809 となり、期待どおり最大の平均合格率が得られていることがわかります。

計算量

クラス数を n とすると、初期化に O(n log n)、以降は追加生徒1人につきヒープ操作が1回発生するため、時間計算量は O(n log n + extra × log n)、空間計算量は O(n) です。単純に毎回全クラスを走査する方法（O(extra × n)）よりも、extra や n が大きいケースで大幅に高速になります。