Pythonで無向グラフの頂点間に低コストのパスが存在するか判定するプログラム

重み付きの無向グラフが与えられているとします。ここで実装が必要なのは、2つの頂点とコストの上限「limit」を引数として受け取り、その上限より低いコストで両頂点を結ぶパスが存在するかどうかを判定する関数 query() です。条件を満たすパスが存在すれば True を、存在しなければ False を返します。

例として、次のようなグラフを考えてみましょう。

このとき、クエリが (0, 2, 10)、(3, 1, 30)、(4, 3, 30) である場合、出力は次のようになります。

False
True
True

出力結果の解説

• 1つ目のクエリ (0, 2, 10) → False: コスト10以内で頂点0から頂点2へ至るパスが存在しないためです。

• 2つ目のクエリ (3, 1, 30) → True: 頂点3から頂点1へはコスト10のパスが存在し、これは制限値の30より小さいためです。

• 3つ目のクエリ (4, 3, 30) → True: 頂点4から頂点3へは、制限値である30以内のコストで到達できるパスが存在するためです。

解法のアプローチ

この問題を解くには、次の手順に従います。

• weights := グラフに含まれる異なる重み(辺のコスト)を格納したリスト

• connections := 各重みの段階における頂点同士の連結状態(所属グループ)を格納したリスト

• 関数 query() を定義します。引数は p、q、limit です。

  • index := リストweightsにおいて、ソート順を保ったままlimitを挿入できる位置(bisect_leftによる二分探索)

  • indexが0の場合(limitより小さい重みの辺が1本も存在しない場合)

    • Falseを返す

  • connections[index-1][p] と connections[index-1][q] が同一であればTrueを返す

この手法のポイントは、辺を重みの昇順にソートし、Union-Find(素集合データ構造)を使って順次連結していくことです。これにより「あるコスト未満の辺のみを使用した場合の連結状態」を各段階で記録できます。クエリ処理時は二分探索で該当する時点の連結情報を参照するだけでよく、大量のクエリにも効率的に対応できます。

実装例

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

import bisect
class Solution(object):

   def __init__(self, n, edgeList):
      def find(node):
         if parent[node]!=node:
            parent[node] = find(parent[node])
         return parent[node]

      def union(x,y):
         parent[find(y)] = find(x)
         return

      parent = {i:i for i in range(n)}
      edgeList.sort(key = lambda x:x[2])
      self.connections = []
      self.weights = []
      for index,(i,j,weight) in enumerate(edgeList):
         union(i,j)
         if index!=len(edgeList)-1 and weight == edgeList[index+1][2]:
            continue
         self.weights.append(weight)
         self.connections.append([find(i) for i in parent])


   def query(self, p, q, limit):
      index = bisect.bisect_left(self.weights,limit)
      if index==0:
         return False
      return self.connections[index-1][p] == self.connections[index-1][q]

ob = Solution(5, [[0, 1, 10], [0, 2, 20], [1, 4, 10], [0, 3, 10], [1, 2, 20], [2, 3, 10]])
print(ob.query(0, 2, 10))
print(ob.query(3, 1, 30))
print(ob.query(4, 3, 30))

入力

ob = Solution(5, [[0, 1, 10], [0, 2, 20], [1, 4, 10], [0, 3, 10], [1, 2, 20], [2, 3, 10]])
print(ob.query(0, 2, 10))
print(ob.query(3, 1, 30))
print(ob.query(4, 3, 30))

出力

False
True
True