Pythonで2人プレイヤーが集められるコインの最大数を求めるアルゴリズム

問題概要

この記事では、2次元マトリックス(グリッド)を対象に、2人のコイン収集者が集められるコインの総数の最大値を求める問題を解説します。マトリックスの各セルにはそのセルにあるコインの枚数が格納されており、1人目の収集者は左上の角、2人目の収集者は右上の角からスタートします。

2人は以下のルールに従って移動します。

• セル (i, j) にいる収集者は、次の行にある (i + 1, j − 1)、(i + 1, j)、(i + 1, j + 1) のいずれかのセルへ移動できます。
• セルに到達した時点で、そのセルのコインをすべて回収し、セルは空になります。
• 収集者は移動せず同じ列にとどまることも可能ですが、各セルのコインを回収できるのは1回だけです。

目標は、この条件下で集められるコインの最大数を見つけることです。

入力例と出力

たとえば、次のようなマトリックスが与えられたとします。

この場合、出力は 17 になります。

解き方:動的計画法によるアプローチ

この問題は、2人の位置を行ごとに追跡する再帰的なDP(動的計画法)で効率よく解けます。手順は以下の通りです。

• A := 入力マトリックス
• R := A の行数
• C := A の列数
• dp(r, c1, c2) 関数を定義する(r は現在の行、c1・c2 はそれぞれ1人目・2人目の現在の列)
  • r == R の場合は、これ以上進めないので 0 を返します。
  • ans = A[r][c1] +(c1 ≠ c2 の場合のみ)A[r][c2] を計算します。2人が同じセルにいる場合、コインは一度しか回収できない点に注意してください。
  • base := ans として保存します。
  • nc1 を [c1 − 1, c1, c1 + 1] の範囲で、nc2 を [c2 − 1, c2, c2 + 1] の範囲で全探索し、両方が有効な列(0 ≤ nc1 < C かつ 0 ≤ nc2 < C)である場合、ans = max(ans, base + dp(r + 1, nc1, nc2)) で更新します。
  • ans を返します。
• 最後に dp(0, 0, C − 1) を呼び出して結果を返します。

実装例(Python)

class Solution:
    def solve(self, A):
        R, C = len(A), len(A[0])
        def dp(r, c1, c2):
            if r == R:
                return 0
            ans = base = A[r][c1] + (c1 != c2) * A[r][c2]
            for nc1 in [c1 − 1, c1, c1 + 1]:
                for nc2 in [c2 − 1, c2, c2 + 1]:
                    if 0 <= nc1 < C and 0 <= nc2 < C:
                        ans = max(ans, base + dp(r + 1, nc1, nc2))
            return ans
        return dp(0, 0, C − 1)
ob = Solution()
print(ob.solve([
    [0, 4, 1, 0],
    [3, 1, 4, 0],
    [2, 5, 1, 1],
    [3, 0, 0, 0]
]))

入力

[
    [0, 4, 1, 0],
    [3, 1, 4, 0],
    [2, 5, 1, 1],
    [3, 0, 0, 0]
]

出力

17

計算量とメモ化について

上記の実装は素朴な再帰のため、入力サイズが大きくなると指数的に処理時間が増加する可能性があります。@lru_cache(functools モジュール)などでメモ化を追加すれば、状態数は「行 × 2人の列の組み合わせ」である O(R × C²) に抑えられ、各状態からの遷移は最大9通りなので、全体の計算量は O(R × C² × 9) となり、実用的な速度で動作します。