Pythonの動的計画法でコンテストの期待得点を最大化するプログラム

複数の問題が出題されるプログラミングコンテストを想定します。このコンテストには少し変わったルールがあり、1問でも正解した時点でコンテストは終了します。

同じ長さを持つ2つのリスト points と chances が与えられます。i番目の問題では、chances[i]% の確率で正解して points[i] ポイントを獲得できます。また、挑戦できる問題数の上限を表す値 k が与えられ、同じ問題に2回以上挑戦することはできません。

この問題のゴールは、最適な戦略を取ったときに獲得できるポイントの期待値を求め、最も近い整数に丸めて返すことです。i番目の問題に挑戦することの期待値は points[i] * chances[i] / 100.0 で計算でき、これは平均的に得られるポイント数を意味します。

たとえば、入力が points = [600, 400, 1000]、chances = [10, 90, 5]、k = 2 の場合、出力は 392 になります。

解き方のアプローチ

この問題は動的計画法（DP）を使って解きます。全体の流れは次のとおりです。

• n := points のサイズとする
• i = 0 ～ n-1 の範囲で、chances[i] := chances[i] / 100.0 とし、パーセントを確率に変換する
• R := インデックス 0 ～ n-1 を、points の降順に並べ替えた順列とする
• int(dp(0, K)) を返す

dp関数の定義

dp(i, k) は「i番目以降の問題だけを対象に、残り k 回の挑戦で得られる期待値」を返します。

• i == n のとき（すべての問題を検討し終えた場合）→ 0.0 を返す
• j := R[i]（配点の高い順に並べたときの i 番目の問題インデックス）
• p := chances[j]
• ev := p * points[j]
• k == 1 のとき → max(ev, dp(i + 1, k)) を返す
• それ以外 → max(dp(i + 1, k - 1) * (1 - p) + ev, dp(i + 1, k)) を返す

漸化式の dp(i + 1, k - 1) * (1 - p) + ev の部分は、「問題jに挑戦して失敗した場合（確率 1-p）は残り k-1 回で先へ進む一方、成功すればその時点でコンテストが終了する」ことを表現しています。もう一方の dp(i + 1, k) は「この問題をスキップする」選択肢に相当します。

実装例

以下のPythonコードで実際の実装を確認してみましょう。再帰呼び出しの重複計算を避けるため、functools.lru_cache によるメモ化を追加しています。

from functools import lru_cache

class Solution:
   def solve(self, points, chances, K):
      n = len(points)
      for i in range(n):
         chances[i] /= 100.0
      R = sorted(range(n), key=points.__getitem__, reverse=True)

      @lru_cache(maxsize=None)
      def dp(i, k):
         if i == n:
            return 0.0
         j = R[i]
         p = chances[j]
         ev = p * points[j]
         if k == 1:
            return max(ev, dp(i + 1, k))
         return max(dp(i + 1, k - 1) * (1 - p) + ev, dp(i + 1, k))

      return int(dp(0, K))

ob = Solution()
print(ob.solve([600, 400, 1000], [10, 90, 5], 2))

入力

[600, 400, 1000], [10, 90, 5], 2

出力

392

出力392の内訳

この結果がどのように導かれるのかを具体的に追ってみましょう。配点の降順に並べると、問題の順序は「1000点（成功率5%）→ 600点（成功率10%）→ 400点（成功率90%）」となります。

まず1000点の問題に挑戦した場合の期待値は 1000 × 0.05 = 50 です。失敗する確率は 95% なので、その場合は残り1回の挑戦で最も期待値の高い選択、つまり 400 × 0.9 = 360 の問題に挑戦することになります。

したがって全体の期待値は 0.05 × 1000 + 0.95 × 360 = 50 + 342 = 392 となり、これが最適な戦略における最大値です。

計算量

メモ化を行うことで、DPの状態数は高々 n × k 個に抑えられ、各状態の遷移は定数時間で処理できます。そのため時間計算量は O(n × k)、空間計算量もメモ化テーブルと再帰スタックを合わせて O(n × k) 程度に収まります。問題数や挑戦回数が増えても現実的な時間で動作する実装になっています。