Union-Find

読了目安約3分

グループ分けを管理し、つながっているかをほぼ O(1) で判定するデータ構造。

この章の目次

「頂点 a と b はつながっているか」という質問に、何度も答える場面を考えます。毎回 DFS で調べると 1 回あたり O(N + M) かかり、質問が多いと間に合いません。

Union-Find は、要素のグループ分けを管理するデータ構造です。次の 2 つの操作を、どちらもほぼ O(1) で行えます。

union：2 つの要素が属するグループを 1 つに合併する
same：2 つの要素が同じグループかを判定する

親をたどって代表を探す

各グループに 1 つ、根と呼ぶ代表の要素を決めます。配列 parent に「自分の親」を記録し、根だけは親を自分自身にします。グループは、根に向かう木の形になります。

親を根までたどる操作が find です。 same は「find の結果が同じか」、union は「片方の根をもう片方の根の子にする」で実現できます。

速くする 2 つの工夫

木が一直線に伸びると、find が O(N) に落ちてしまいます。そこで 2 つの工夫を入れます。

経路圧縮：find でたどった要素を、すべて根に直接つなぎ直す
サイズによる合併：union では要素数の少ない木を多い木の下につなぐ

この 2 つを入れると、木がほぼ平らに保たれ、1 回の操作はほぼ O(1) になります。

正確な計算量は「アッカーマン関数の逆関数」に比例します。実用上 4 以下にしかならない値なので、定数時間とみなして差し支えありません。

struct で実装する

Swift

struct UnionFind {
    private var parent: [Int]
    private var size: [Int]

    init(count: Int) {
        parent = Array(0..<count)
        size = [Int](repeating: 1, count: count)
    }

    mutating func find(_ x: Int) -> Int {
        if parent[x] == x { return x }
        parent[x] = find(parent[x])  // 経路圧縮
        return parent[x]
    }

    mutating func union(_ a: Int, _ b: Int) {
        var rootA = find(a)
        var rootB = find(b)
        if rootA == rootB { return }
        if size[rootA] < size[rootB] { swap(&rootA, &rootB) }
        parent[rootB] = rootA  // サイズによる合併
        size[rootA] += size[rootB]
    }

    mutating func same(_ a: Int, _ b: Int) -> Bool {
        find(a) == find(b)
    }
}

var uf = UnionFind(count: 5)
uf.union(0, 1)
uf.union(2, 3)
print(uf.same(0, 1))
print(uf.same(1, 2))
uf.union(1, 2)
print(uf.same(0, 3))

初期状態は「全員が 1 人グループ」で、parent は自分自身、size は 1 です。 find が parent を書き換えるため、各メソッドに mutating を付けています。

試してみよう: 最後に print(uf.same(0, 4)) を足して、union していない要素 4 だけが別グループのままなことを確かめてください。

学びどころ

概念	一言まとめ
`parent` 配列	親を記録し、根がグループの代表になる
find	根までたどる。経路圧縮で木を平らにする
union	小さい木を大きい木の根につなぐ
same	find の結果を比べる。ほぼ O(1)

演習

1 行目に要素数 N とクエリ数 Q が空白区切りで入力されます。続く Q 行に、クエリ t a b (0 ≦ a, b ≦ N - 1) が入力されます。

t が 0 なら、a と b の属するグループを合併する
t が 1 なら、a と b が同じグループなら Yes、違えば No を出力する

テキスト
入力例:
4 5
1 0 1
0 0 1
1 0 1
0 2 3
1 1 2

出力例:
No
Yes
No

模範解答

Swift
struct UnionFind {
    private var parent: [Int]
    private var size: [Int]

    init(count: Int) {
        parent = Array(0..<count)
        size = [Int](repeating: 1, count: count)
    }

    mutating func find(_ x: Int) -> Int {
        if parent[x] == x { return x }
        parent[x] = find(parent[x])
        return parent[x]
    }

    mutating func union(_ a: Int, _ b: Int) {
        var rootA = find(a)
        var rootB = find(b)
        if rootA == rootB { return }
        if size[rootA] < size[rootB] { swap(&rootA, &rootB) }
        parent[rootB] = rootA
        size[rootA] += size[rootB]
    }

    mutating func same(_ a: Int, _ b: Int) -> Bool {
        find(a) == find(b)
    }
}

let nq = readLine()!.split(separator: " ").map { Int($0)! }
let n = nq[0]
let q = nq[1]
var uf = UnionFind(count: n)
for _ in 0..<q {
    let query = readLine()!.split(separator: " ").map { Int($0)! }
    if query[0] == 0 {
        uf.union(query[1], query[2])
    } else {
        print(uf.same(query[1], query[2]) ? "Yes" : "No")
    }
}