Pythonでプリムのアルゴリズムを使って最小全域木（MST）を求める方法

最小全域木（MST）とは？

グラフが与えられたとき、そこから「最小全域木」（MST：Minimum Spanning Tree）を求めることを考えます。グラフのMSTとは、重み付きグラフの部分集合であり、すべての頂点が含まれており互いに接続され、かつ部分集合内に閉路（サイクル）が存在しないものを指します。「最小」と呼ばれるのは、MSTの辺の重みの合計が、元のグラフから構成できるどの全域木よりも小さくなるためです。

この記事では、プリム（Prim）のMSTアルゴリズムを実装し、与えられたグラフからMSTの辺の重みの合計を求める方法を解説します。

問題の例

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

このとき、頂点数 n = 4、開始頂点 s = 3 である場合、出力は 14 になります。

このグラフから得られるMSTは次のようになります。

このMSTの辺の重みの合計は 14 です。

解き方の手順

この問題を解くために、以下の手順に従います。

• mst_find() 関数を定義する。引数は G（グラフ）と s（開始頂点）。
  • distance := グラフGと同じサイズのリストを無限大（inf）で初期化
  • distance[s] := 0
  • itr := グラフGと同じサイズのリストをFalseで初期化（訪問済みフラグ）
  • c := 0（MSTの重み合計）
  • Trueの間、以下を繰り返す：
    • min_weight := 無限大
    • m_idx := -1
    • i を 0 から Gのサイズまで繰り返す：
      • itr[i] が False の場合：
        • distance[i] < min_weight であれば：
          • min_weight := distance[i]
          • m_idx := i
    • m_idx が -1 の場合：
      • ループを抜ける
    • c := c + min_weight
    • itr[m_idx] := True
    • G[m_idx] 内の各ペア (i, j) について：
      • distance[i] := distance[i] と j の最小値
  • c を返す
• G := 頂点数n個の空の辞書（隣接リスト）を持つ辞書を作成
• edges の各要素 item について：
  • u := item[0]
  • v := item[1]
  • w := item[2]
  • u := u - 1（0始まりのインデックスに変換）
  • v := v - 1
  • 既存の辺がある場合は重みの小さい方を採用し、G[u][v] と G[v][u] に設定（無向グラフとして両方向に登録）
• mst_find(G, s) を返す

Pythonでの実装例

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

def mst_find(G, s):
    distance = [float("inf")] * len(G)
    distance[s] = 0
    itr = [False] * len(G)
    c = 0
    while True:
        min_weight = float("inf")
        m_idx = -1
        for i in range(len(G)):
            if itr[i] == False:
                if distance[i] < min_weight:
                    min_weight = distance[i]
                    m_idx = i
        if m_idx == -1:
            break
        c += min_weight
        itr[m_idx] = True
        for i, j in G[m_idx].items():
            distance[i] = min(distance[i], j)
    return c

def solve(n, edges, s):
    G = {i: {} for i in range(n)}
    for item in edges:
        u = item[0]
        v = item[1]
        w = item[2]
        u -= 1
        v -= 1
        try:
            min_weight = min(G[u][v], w)
            G[u][v] = min_weight
            G[v][u] = min_weight
        except KeyError:
            G[u][v] = w
            G[v][u] = w
    return mst_find(G, s)

print(solve(4, [(1, 2, 5), (1, 3, 5), (2, 3, 7), (1, 4, 4)], 3))

コードのポイント

• distance配列：各頂点をMSTに追加する際のコスト（未訪問の頂点への最小辺の重み）を保持します。
• 貪欲法（Greedy）：毎回、未訪問の頂点の中から最もコストの小さい頂点を選んでMSTに追加していきます。
• 計算量：この実装では頂点数をVとしたとき O(V²) の計算量になります。優先度付きキュー（heapq）を使うことで O(E log V) まで高速化できます。

入力

4, [(1, 2, 5), (1, 3, 5), (2, 3, 7), (1, 4, 4)], 3

出力

14