Pythonで家の塗装にかかる最小コストを求めるプログラム

問題の概要

小さな町に m 軒の家があるとします。各家庭は n 色（1〜n のラベル付き）のうちどれか一色で塗らなければなりませんが、すでに塗装済みの家もあり、その場合は塗り直す必要はありません。

同じ色で塗られた連続する家のまとまりを「街区（neighborhood）」と呼びます。入力として与えられるデータは次のとおりです。

• houses[i]：i 番目の家の色。値が 0 の場合はまだ塗装されていないことを表します。
• costs[i][j]：i 番目の家を色 j+1 で塗るときのコスト（2次元配列）。
• target：最終的に作りたい街区の数。

求めるのは、残りの家をすべて塗装した結果、街区の数がちょうど target 個になるときの最小コストです。条件を満たす塗り方が存在しない場合は -1 を返します。

入出力の例

たとえば、次の入力を考えてみましょう。

houses = [0,2,1,2,0]
cost = [[1,10],[10,1],[10,1],[1,10],[5,1]]
n = 2
target = 3

この場合の出力は 11 になります。すでに塗装済みの家もあるため、残りを [2,2,1,2,2] のように塗れば、街区は {2,2}、{1}、{2,2} の 3 つに分かれます。このとき最初と最後の家を塗るコストは 10 + 1 = 11 となり、これが最小です。

解法のアプローチ

この問題は再帰を使って次の手順で解きます。

• m := houses のサイズとします。
• 関数 helper(i, p_col, grp) を定義します。i は現在見ている家の番号、p_col は直前の家の色、grp はそこまでの街区数です。
• i == m のとき（すべての家を見終わったとき）、grp == target なら 0、そうでなければ無限大（inf）を返します。
• houses[i] != 0 のとき（すでに塗装済み）、helper(i + 1, houses[i], grp + (p_col != houses[i] なら 1、そうでなければ 0)) を返します。
• total := inf とし、col を 1 から n まで動かしながら、total = min(total, cost[i][col-1] + helper(i + 1, col, grp + (p_col != col なら 1、そうでなければ 0))) を計算します。
• total を返します。
• メイン処理では ans := helper(0, -1, 0) を呼び出し、ans が inf でなければ ans を、inf なら -1 を返します。

Pythonでの実装例

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

def solve(houses, cost, n, target):
   m = len(houses)

   def helper(i, p_col, grp):
      if i == m:
         return 0 if grp == target else float('inf')

      if houses[i] != 0:
         return helper(i + 1, houses[i], grp + int(p_col != houses[i]))

      total = float('inf')
      for col in range(1, n + 1):
         total = min(total, cost[i][col - 1] + helper(i + 1, col, grp + int(p_col != col)))

      return total

   ans = helper(0, -1, 0)
   return ans if ans != float('inf') else -1

houses = [0,2,1,2,0]
cost = [[1,10],[10,1],[10,1],[1,10],[5,1]]
n = 2
target = 3

print(solve(houses, cost, n, target))

実行結果

入力：

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

出力：

11

補足：計算量と改善のポイント

上記の素朴な再帰は、各未塗装の家について n 通りの色を試すため、最悪の場合 O(m^n) の指数時間がかかります。実務では functools.lru_cache などでメモ化を行うか、「家の番号 × 直前の色 × 街区数」を状態とする動的計画法（DP）に書き換えることで、O(m × n² × target) 程度まで計算量を抑えられます。特に m や n が大きいケースでは、メモ化の導入がパフォーマンス向上の鍵となります。