Pythonで全アイテムをちょうど1つずつ購入する最小コストを求めるプログラム

N個のアイテムがあり、それぞれ0、1、2、…、N-1という番号が付けられているとします。ここで、サイズSの2次元リストsetsが与えられます。i番目のセットは価格sets[i][2]で購入することができ、sets[i][0]からsets[i][1]までの範囲にあるすべてのアイテムを受け取れます。さらに、サイズNのリストremovalsも与えられ、i番目の要素のインスタンスを1つ、価格removals[i]で廃棄できます。

このとき、0からN-1までの各要素をちょうど1つずつ入手するための最小コストを求めてください。どうしても実現できない場合は-1を返します。

入力例

sets = [
    [0, 4, 4],
    [0, 5, 12],
    [2, 6, 9],
    [4, 8, 10]
]
removals = [2, 5, 4, 6, 8]

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

解法のアプローチ

この問題は、グラフの最短経路問題としてモデル化できます。各セット購入を「前進する辺」、重複したアイテムの廃棄を「後退する辺」とみなし、ダイクストラ法(Dijkstra法)を用いてノード0からノードNへの最短距離を求めます。具体的な手順は以下の通りです。

• N := removals のサイズとする

• graph := サイズ (N + 1) × (N + 1) の新しい隣接リストを作成する

• sets 内の各要素 s, e, w について:

  • graph[s] に [e+1, w] を追加する

• removals 内の各インデックス i と値 r について:

  • graph[i + 1] に [i, r] を追加する

• pq := 新しい優先度付きキュー(ヒープ)を作成する

• dist := 新しい辞書(距離マップ)を作成する

• dist[0] := 0 と初期化する

• pq が空でない限り、以下を繰り返す:

  • d, e := ヒープ pq から最小のコストを持つ項目を取り出す

  • もし dist[e] < d が成り立つなら(古い情報の場合):

    • 次の反復へスキップする

  • もし e が N と等しいなら:

    • d を返す(ゴール到達)

  • graph[e] 内の各隣接ノード nei と重み w について:

    • d2 := d + w を計算する

    • もし d2 < dist[nei] が成り立つなら:

      • dist[nei] := d2 と更新する

      • [d2, nei] を pq に追加する

• ループが終了したら -1 を返す(到達不可能)

実装例

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

import heapq
from collections import defaultdict
class Solution:
   def solve(self, sets, removals):
      N = len(removals)
      graph = [[] for _ in range(N + 1)]
      for s, e, w in sets:
         graph[s].append([e + 1, w])
      for i, r in enumerate(removals):
         graph[i + 1].append([i, r])
      pq = [[0, 0]]
      dist = defaultdict(lambda: float("inf"))
      dist[0] = 0
      while pq:
         d, e = heapq.heappop(pq)
         if dist[e] < d:
            continue
         if e == N:
            return d
         for nei, w in graph[e]:
            d2 = d + w
            if d2 < dist[nei]:
               dist[nei] = d2
               heapq.heappush(pq, [d2, nei])
      return -1

ob = Solution()
print (ob.solve([
   [0, 4, 4],
   [0, 5, 12],
   [2, 6, 9],
   [4, 8, 10]
], [2, 5, 4, 6, 8]))

入力

[[0, 4, 4],
[0, 5, 12],
[2, 6, 9],
[4, 8, 10]], [2, 5, 4, 6, 8]

出力

4