Pythonで都市を最小コストで接続する方法｜クラスカル法とUnion-Findによる実装

問題概要

1からNまでの番号が付けられたN個の都市があるとします。接続情報connectionsの各要素は[city1, city2, cost]という形式で与えられ、これはcity1とcity2を直接つなぐためのコストを表します。ここで求めたいのは、任意の2つの都市の間に必ず経路が存在する状態（全域木）を作るときの最小コストです。コストは採用した接続のコストの合計であり、すべての都市を接続できない場合は-1を返します。

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

この場合の出力は6になります。3つの都市をすべてつなぐには2本の接続で十分なので、コストの小さい組み合わせ、すなわち[2, 3, 1]と[1, 2, 5]を採用すれば、合計コストは1 + 5 = 6となります。

解き方のアプローチ

この問題は、クラスカル法（Kruskal's algorithm）と呼ばれる最小全域木（MST）のアルゴリズムを使うことで効率的に解けます。具体的には、以下の手順に従います。

• find()メソッドを定義します。引数はxです。

• parent[x]が-1であれば、xをそのまま返します。

• そうでなければparent[x] := find(parent[x])として再帰的に根を辿ります（経路圧縮）。

• parent[x]を返します。

• union()メソッドを定義します。引数はxとyです。

• parent_x := find(x)、parent_y := find(y)を求めます。

• parent_xとparent_yが等しければ何もしません（すでに同じグループ）。異なればparent[parent_y] := parent_xとしてグループを統合します。

• メインのメソッドではnとconnectionsを受け取ります。

• parentをサイズn+1の配列として-1で初期化し、disjoint := n - 1、cost := 0とします。

• c := connectionsをコスト（3番目の要素）を基準に昇順ソートしたリストを作ります。

• i := 0 とします。

• i < len(c) かつ disjoint が0でない間、次を繰り返します。

  • x := c[i][0]、y := c[i][1] とします。
  • find(x)とfind(y)が異なる場合のみ、disjointを1減らし、costにc[i][2]を加算してunion(x, y)を実行します。
  • i を1増やします。

• disjointが0になっていればcostを返し、そうでなければ-1を返します。

アルゴリズムのポイント

コストの安い接続から順に検討し、まだつながっていない都市同士をつなげていくのがクラスカル法の基本的な考え方です。「2つの都市がすでにつながっているかどうか」の判定には、Union-Find（素集合データ構造）を使用します。さらに、find()の中で親を付け替える「経路圧縮」を行うことで、以降の探索が高速化されます。

Pythonでの実装例

それでは、実際のPythonコードを見てみましょう。

class Solution(object):
   def find(self, x):
      if self.parent[x] == -1:
         return x
      self.parent[x] = self.find(self.parent[x])
      return self.parent[x]
   def union(self, x, y):
      parent_x = self.find(x)
      parent_y = self.find(y)
      if parent_x == parent_y:
         return
      self.parent[parent_y] = parent_x
   def minimumCost(self, n, connections):
      self.parent = [-1 for i in range(n + 1)]
      disjoint = n - 1
      cost = 0
      c = sorted(connections, key=lambda v: v[2])
      i = 0
      while i < len(c) and disjoint:
         x = c[i][0]
         y = c[i][1]
         if self.find(x) != self.find(y):
            disjoint -= 1
            cost += c[i][2]
            self.union(x, y)
         i += 1
      return cost if not disjoint else -1

ob = Solution()
print(ob.minimumCost(3, [[1,2,5], [1,3,6], [2,3,1]]))

入力

3
[[1,2,5],[1,3,6],[2,3,1]]

出力

6

実行結果の解説

まずコスト1の接続[2, 3]を採用し、次にコスト5の接続[1, 2]を採用すると、disjointが0になり、3つの都市すべてが接続されます。このときの合計コストは1 + 5 = 6です。仮に接続が不足していてすべての都市をつなげられない場合は、disjointが0にならないままループが終了するため、-1が返されます。

計算量について

接続リストのソートにO(E log E)（Eは接続の本数）、Union-Findの各操作は経路圧縮によりほぼ定数時間で処理できるため、全体の計算量はO(E log E)となります。また、必要な接続は最大N-1本であるため、disjointが0になった時点でループを打ち切れる点も効率的です。