Pythonで加重グラフの最小コストを求めるプログラムの実装方法

問題の概要

整数の2次元リスト edges が与えられます。これは無向グラフを表しており、各行は1本の辺 [u, v, w] に対応します。つまり、ノード u とノード v が接続されており、その辺の重みが w であることを意味します。グラフは 0 から n-1 までの n 個のノードで構成されています。

ここで「パスのコスト」は、パスに含まれる辺の数とパス上の辺の重みの最大値の積として定義されます。求めるのは、ノード 0 からノード n-1 へ到達するときの最小コストであり、そのようなパスが存在しない場合は -1 を返します。

たとえば、入力が次のようなケースを考えてみましょう。

edges = [
    [0, 2, 100],
    [1, 2, 200],
    [1, 3, 100],
    [2, 3, 300]
]

このとき出力は 600 になります。パス「0 → 2 → 1 → 3」は3本の辺からなり、その最大重みは 200 であるため、コストは 3 × 200 = 600 となるからです。

解き方のアプローチ

この問題は、「使用できる辺の重みの上限(weight_cap)を段階的に下げながら幅優先探索(BFS)を繰り返す」という戦略で解けます。上限を下げるほど通れる辺が減り、見つかるパスの最大重みも小さくなるため、最終的に最小コストが得られます。具体的な手順は以下の通りです。

ステップ1: グラフデータの構築

• graph: 隣接ノードを格納するマップ(defaultdict)
• weights: 辺の重みを格納するマップ
• max_weight := 0、N := 0 で初期化
• edges 内の各 (u, v, w) について:
  • graph[u] の末尾に v を追加し、graph[v] の末尾に u を追加(無向グラフのため両方向)
  • weights[(u, v)] と weights[(v, u)] に w を設定
  • N を max(N, u+1, v+1) で更新
  • max_weight を max(max_weight, w) で更新

ステップ2: 重み上限を下げながらBFSを繰り返す

• result := 無限大 で初期化
• max_weight >= 0 の間、以下を繰り返します:
  • d, weight := bfs(0, max_weight) を呼び出す
  • d >= 0(パスが見つかった)場合:
    • result を min(result, d × weight) で更新
    • max_weight を weight − 1 に下げて再挑戦
  • それ以外(パスが存在しない)場合はループを終了
• 最後に、result が無限大より小さければ result を、そうでなければ -1 を返します

bfs() 関数の内部処理

bfs(root, weight_cap) は、「重みが weight_cap 以下の辺のみを使用して」ノード root からノード N-1 へ至る最短パス(辺数が最小のパス)を探索します。

• キュー Q を deque として用意し、初期値 (root, 0, 0)(ノード、距離、現在の最大重み)を設定
• 訪問管理用の配列 visited を作成し、visited[0] を True に設定
• Q が空でない間、以下を繰り返します:
  • Q の末尾から要素 (v, d, current_weight) を取り出す
  • v が N-1 なら、(d, current_weight) を返す
  • graph[v] 内の各隣接ノード w について:
    • visited[w] が True ならスキップ
    • new_weight := weights[(v, w)] とし、new_weight <= weight_cap であれば:
      • visited[w] := True とする
      • (w, d+1, max(current_weight, new_weight)) を Q の左端に追加
• パスが見つからなければ (-1, -1) を返します

Pythonでの実装例

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

from collections import defaultdict, deque


class Solution:
    def solve(self, edges):
        graph = defaultdict(list)
        weights = {}
        max_weight = 0
        N = 0
        for u, v, w in edges:
            graph[u].append(v)
            graph[v].append(u)
            weights[u, v] = w
            weights[v, u] = w
            N = max(N, u + 1, v + 1)
            max_weight = max(max_weight, w)

        def bfs(root, weight_cap):
            Q = deque([(root, 0, 0)])
            visited = [False] * N
            visited[0] = True
            while Q:
                v, d, current_weight = Q.pop()
                if v == N - 1:
                    return d, current_weight
                for w in graph[v]:
                    if visited[w]:
                        continue
                    new_weight = weights[v, w]
                    if new_weight <= weight_cap:
                        visited[w] = True
                        Q.appendleft((w, d + 1, max(current_weight, new_weight)))
            return -1, -1

        result = float("inf")
        while max_weight >= 0:
            d, weight = bfs(0, max_weight)
            if d >= 0:
                result = min(result, d * weight)
                max_weight = weight - 1
            else:
                break
        return result if result < float("inf") else -1


ob = Solution()
print(ob.solve([
    [0, 2, 100],
    [1, 2, 200],
    [1, 3, 100],
    [2, 3, 300]
]))

入力

[
    [0, 2, 100],
    [1, 2, 200],
    [1, 3, 100],
    [2, 3, 300]
]

出力

600

計算量の目安

BFS 1回あたりの計算量は O(V + E) です。メインループでは、パスが見つかるたびに重みの上限が「そのパスで使われた最大重み − 1」に更新されるため、ループの反復回数はグラフに現れる重みの種類数で抑えられます。したがって、全体の計算量は O(E × (V + E)) 程度と見積もられます。