Pythonでエッジの重みに上限があるパスの存在を判定するプログラム

問題概要

n 個のノードを持つ無向重み付きグラフを考えます。グラフは edgeList で与えられ、edgeList[i] は (u, v, w) の3つの要素からなり、「u と v の間に距離 w の辺が存在する」ことを表します。

さらに、query[i] が (p, q, lim) の形式を持つクエリ配列も与えられます。各クエリは「p から q へ(直接または他のノードを経由して)、距離が lim 未満の経路は存在するか?」を問うものです。すべてのクエリに対する True / False の結果を配列として返す必要があります。

具体例

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

このとき出力は [True, False, True] になります。

• 1 → 4:経路 1 → 3 → 4 をたどればコスト 11 で到達できるため True
• 2 → 3:コスト 3 未満では移動できないため False
• 1 → 2:経路 1 → 3 → 2 をたどればコスト 14 で、これは 15 未満なので True

解法のアプローチ:Union-Findによるオフライン処理

この問題は Union-Find(素集合データ構造) を使うことで効率的に解けます。ポイントは次のとおりです。

• 辺を重みの昇順にソートする
• クエリを制限値 lim の昇順にソートする
• クエリを小さい lim から順に処理しながら、「その時点までの lim 未満の辺」だけを順次マージしていく
• 各クエリの時点で p と q が同じ連結成分に属しているかを find() で判定する

こうすると、各クエリについて「lim 未満の辺のみを使って到達可能か」を正しく判定でき、計算量は O(E log E + Q log Q + (E + Q) α(n)) に抑えられます(α はアッカーマン逆関数で、ほぼ定数扱いできます)。

アルゴリズムの手順

• parent := 0 から n までの連番リストを作成
• rank := サイズ n+1 のリストを 0 で初期化
• find() 関数を定義(parent と x を引数に取る)
  • parent[x] が x と等しければ x を返す(根に到達)
  • そうでなければ parent[x] := find(parent, parent[x]) で経路圧縮し、parent[x] を返す
• union() 関数を定義(parent、a、b を引数に取る)
  • a := find(parent, a)、b := find(parent, b)
  • a == b なら何もしない(すでに同じ集合)
  • rank[a] < rank[b] なら parent[a] := b
  • rank[a] > rank[b] なら parent[b] := a
  • それ以外の場合は parent[b] := a とし、rank[a] を 1 増やす(ランクによる平衡化)
• メイン処理:
  • edgeList を重み(第3要素)を基準にソート
  • res := クエリ数と同じ長さの結果配列
  • queries にインデックス i を付けて (i, クエリ内容) のペアのリストにし、lim(第3要素)を基準にソート
  • ind := 0 として、各クエリ (a, b, w) について:
    • ind が edgeList の範囲内かつ edgeList[ind][2] < w の間、union() で辺をマージし ind を進める
    • res[i] := (find(parent, a) == find(parent, b)) を設定
  • res を返す

実装例(Python)

以下のコードで動作を確認できます。

def solve(n, edgeList, queries):
    parent = [i for i in range(n+1)]
    rank = [0 for i in range(n+1)]

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

    def union(parent, a, b):
        a = find(parent, a)
        b = find(parent, b)
        if a == b:
            return
        if rank[a] < rank[b]:
            parent[a] = b
        elif rank[a] > rank[b]:
            parent[b] = a
        else:
            parent[b] = a
            rank[a] += 1

    edgeList.sort(key=lambda x: x[2])
    res = [False] * len(queries)
    queries = [[i, ch] for i, ch in enumerate(queries)]
    queries.sort(key=lambda x: x[1][2])

    ind = 0
    for i, (a, b, w) in queries:
        while ind < len(edgeList) and edgeList[ind][2] < w:
            union(parent, edgeList[ind][0], edgeList[ind][1])
            ind += 1
        res[i] = find(parent, a) == find(parent, b)
    return res

n = 4
edgeList = [(1,2,16),(1,3,8),(2,4,3),(2,3,6),(4,3,3)]
queries = [(1,4,12),(2,3,3),(1,2,15)]
print(solve(n, edgeList, queries))

入力

4, [(1,2,16),(1,3,8),(2,4,3),(2,3,6),(4,3,3)],[(1,4,12),(2,3,3),(1,2,15)]

出力

[True, False, True]

まとめ

辺とクエリをそれぞれソートして Union-Find でオフライン処理する手法により、各クエリごとにグラフを再構築することなく、「距離 lim 未満で到達可能か」という判定をまとめて高速に行うことができます。経路圧縮とランクによるマージを組み合わせることで、大規模なグラフ・大量のクエリにも十分対応できる実装になっています。