Pythonで隣接するフェンスが同じ色にならないようK色で塗る最小コストを求めるプログラム

問題の概要

N本のフェンスを一列に並べ、K種類の異なる色で塗ることを考えます。ただし「隣り合うフェンス同士は同じ色にできない」という制約があり、そのうえで総コストを最小化したいのです。

入力として N × K の行列が与えられます。n 行 k 列目の値は「n 番目のフェンスを k 番目の色で塗るときのコスト」を表します。このとき、制約を満たす塗り方の中で最小となる総コストを求めます。

たとえば、次のような入力が与えられたとします。

この場合の出力は 14 になります。最初のフェンスから順に、コスト 5 → 2 → 3 → 4 の色を選べば、隣接するフェンスの色が重ならず、合計コストを 14 に抑えられるからです。

解き方（アルゴリズム）

この問題は動的計画法で効率よく解けます。ポイントは、前の行（直前のフェンス）について「最小コスト」と「二番目に小さいコスト」、およびそれぞれに対応する色のインデックスを保持しておくことです。こうすれば、現在のフェンスでどの色を選んでも、直前のフェンスと異なる色なら最小コストを、同じ色なら二番目の最小コストを即座に参照できます。

具体的には、次の手順で進めます。

• n := 行列の行数とします
• fc := 0、ft := 0（fc：前回の最小コストの色インデックス、ft：そのコスト）
• sc := 1、st := 0（sc：前回の二番目の色インデックス、st：そのコスト）
• 行列の各行に対して以下を実行します
  • nfc := -1、nft := 無限大
  • nsc := -1、nst := 無限大
  • 行内の各インデックス i と値 t について
    • ct := t +（i が fc と異なる場合は ft、同じ場合は st）
    • ct が nft 以下の場合：
      nsc, nst := nfc, nft を退避し、nfc, nft := i, ct と更新
    • そうでなく ct が nst 以下の場合：
      nsc, nst := i, ct と更新
  • fc, ft := nfc, nft、sc, st := nsc, nst と更新
• 最後に ft を返します

各行の処理は K 色を一度走査するだけなので、全体の計算量は O(N × K)、追加メモリは O(1) で済みます。すべての色の組み合わせを試す O(N × K²) の素朴なアプローチよりも大幅に高速です。

実装例

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

class Solution:
   def solve(self, matrix):
      n = len(matrix)
      fc, ft = 0, 0
      sc, st = 1, 0
      inf = int(1e18)
      for row in matrix:
         nfc, nft = -1, inf
         nsc, nst = -1, inf
         for i, t in enumerate(row):
            ct = t + (ft if i != fc else st)
            if ct <= nft:
               nsc, nst = nfc, nft
               nfc, nft = i, ct
            elif ct <= nst:
               nsc, nst = i, ct
         fc, ft = nfc, nft
         sc, st = nsc, nst
      return ft

ob = Solution()
matrix = [
   [6, 4, 5],
   [3, 2, 7],
   [3, 4, 5],
   [5, 4, 4]
]
print(ob.solve(matrix))

入力

[
   [6, 4, 5],
   [3, 2, 7],
   [3, 4, 5],
   [5, 4, 4]
]

出力

14