【図で解説】０からわかるUnion-Find入門【C++】

そまちょブログのそまちょ（@somachob）です。

この記事では、データ構造の1つである Union-Find について解説します。

Union-Find とは

Union-Find は、グループ分けを効率的に管理できるデータ構造です。

Union-Find を使うことで、次のようなことが簡単にできるようになります。

要素 x と要素 y のグループを併合する（ Union ）
2つの要素 x, y が同じグループに属するか調べる（ Find ）

Union-Find を使えば、2つの要素が同じグループか調べたり、グループの併合を簡単に行うことができます。

Union-Find の基本的な実装コード

この記事の最終的な Union-Find のコードを紹介します。

struct UnionFind {
    // 親の要素とサイズを管理する変数
    vector<int> parent, size;

    // 変数の初期化
    UnionFind(int n) : parent(n, -1), size(n, 1) {}

    // x の根を求める
    int root(int x) {
        // x が根のとき
        if(parent[x] == -1) return x;

        // 経路圧縮
        return parent[x] = root(parent[x]);
    }

    // x と y の根が同じか
    bool isSame(int x, int y) {
        return root(x) == root(y);
    }

    // x と y のグループを併合する
    void unite(int x, int y) {
        // x と y の根を取得
        int rootX = root(x);
        int rootY = root(y);

        // x と y が同じグループのときは何もしない
        if (rootX == rootY) return;

        // union by size（ y のサイズが小さくなるように調整 ）
        if (size[rootX] < size[rootY]) swap(rootX, rootY);

        // y の親を x にする
        parent[rootY] = rootX;

        // x のサイズに y のサイズを足す
        size[rootX] += size[rootY]; 
    }
};

Union-Find について、これから詳しく解説していきます。

まずは考え方について説明し、そのあとにシンプルな実装を紹介します。その後、計算量を少なくするための union by size と経路圧縮を実装して上のコードになります。

基本的な考え方

７人でグループ分けをすることになりました。各グループには、１人リーダーがいます。

グループ分けをした結果、次の３つのグループができました。それぞれのグループのリーダーは ②、③、④ です。

①のリーダーは②で、⑤のリーダーも②です。

⓪のリーダーは②ですが、⑥のリーダーは④です。

リーダーが同じとき、その人たちは同じグループ。リーダーが違うとき、その人たちは違うグループということになります。

次に、③がリーダーのグループと⑥の所属するグループが併合しました。これにより、④と⑥のリーダーは③に変わりました。

併合するときは、併合するグループの各要素について、新しいリーダーを設定すればいいということになります。

リーダーが同じなら同じグループ
リーダーが別なら違うグループ
グループが併合したら、どちらかのグループのメンバーのリーダが変わる

Union-Find でのデータの扱い方

Union-Find では、グループの扱いを根付き木で表現します。

先ほどのグループ分けを根付き木として表現すると、たとえば次のようになります。

リーダーのことを根付き木では、根といいます。

同じグループかの判定とグループの併合

要素 x の根を root(x) とするとき、要素 x と要素 y が同じグループであるかは次のように判定できます。

root(x) == root(y) ：同じグループ
root(x) != root(y) ：違うグループ

要素 x と要素 y を併合するときは、root(x) の親を root(y) になるようにつなぎます。

次の図は、要素 1 と要素 4 の併合を表しています。

要素 1 の根である要素 2 の親を、要素 3 になるようにつないでいます。

簡単な Union-Find の実装

Union-Find の実装方法ついて説明していきます。

※わかりやすく説明するために、この記事の初めに紹介した Union-Find のコードとは、ところどころ違います。

根を求める root() 関数

Union-Find で大切なのは、要素 x の根を表す root(x) です。

同じグループか判定するときも、グループを併合するときも、根を使います。

要素 x の根を求めるにはどうすればいいでしょうか？

これは、配列 parent を使えば解決できます。

各要素について、自分の親を配列 parent で管理します。こうすることで、ある要素から順々に親をたどっていけば、いつかは根にたどり着きます。

次の図は、２つの根付き木のグループに分かれているときの、０から６までの各要素の配列 parent を表しています。

これをふまえて、要素 x の根を求める root() 関数のコードは次のようになります。配列 parent で、根は -1 で表します。

int root(int x) {
    // x が根のとき
    if(parent[x] == -1) return x;

    // 再帰呼び出し（ x の親を引数に ）
    return root(parent[x]);
}

このコードでは、根を再帰的に求めています（６行目の return root(parent[x]); のところ）。

たとえば、配列 parent が次の値だったとき、要素 0 の根を求めるために root(0) を呼び出すとどうなるでしょうか。

vector<int> parent = {5, 2, -1, -1, 3, 2, 4};

root(0) が呼び出される。parent[0] は 5 のため６行目は return root(5) になる。
root(5) が呼び出される。parent[5] は 2 のため６行目は return root(2) になる。
root(2) が呼び出される。parent[2] は -1 のため３行目の return x の部分が実行される。
x は 2 のため root(2) は 2 を返す。
return root(2) は 2 のため root(5) は 2 を返す。
return root(5) は 2 のため root(0) は 2 を返す。

このように、root(0) を呼び出すと 2 が返されます。

コードを実行して確認してみます。

#include <bits/stdc++.h>
using namespace std;

// 配列 parent
vector<int> parent = {5, 2, -1, -1, 3, 2, 4};

int root(int x) {
    // 説明用の出力
    cout << "root(" << x << ")が呼び出されました。" << endl;
    cout << "parent[" << x << "]は" << parent[x] << "です" << endl;
    cout << endl;

    // x が根のとき
    if(parent[x] == -1) return x;

    // 再帰呼び出し（ x の親を引数に ）
    return root(parent[x]);
}

int main () {
    // root(0) を呼び出して根を取得
    int answer = root(0);

    cout << "root(0)は"<< answer << "です。" << endl;
    return 0;
}

// 実行結果
/*
root(0)が呼び出されました。
parent[0]は5です

root(5)が呼び出されました。
parent[5]は2です

root(2)が呼び出されました。
parent[2]は-1です

root(0)は2です。
*/

2つの要素が同じグループかを判定する isSame() 関数

root() 関数で、要素 x の根が求めれました。

２つの要素の根が同じかどうかを判定すれば、グループが同じかわかります。

同じグループかを判定する isSame() 関数は次のように実装できます。

bool isSame(int x, int y) {
    return root(x) == root(y);
}

このコードがわかりにくければ、次のコードを見てください。していることは同じです。

bool isSame(int x, int y) {
    int rootX = root(x);  // => x の根を取得してrootXに代入
    int rootY = root(y);  // => y の根を取得してrootYに代入

    // rootX と rootY が同じかどうか判定して分岐
    if (rootX == rootY) return true;
    else return false;
}

2つのグループを併合する unite() 関数

要素 x と要素 y を併合する関数について説明します。

グループを併合するためには、併合する２つの要素のグループは別グループである必要があります。

そのため、2つの要素の根が同じときは、すでに同じグループなので併合する必要はありません。

グループを併合するときは、片方のグループの根の親を、もう片方のグループの根になるように設定します。

これをふまえて、実装した unite() 関数が次のコードです。

void unite(int x, int y) {
    // x と y が同じグループのときは何もしない
    if (isSame(x, y)) return;

    // x の根の親を y の根にする
    parent[root(x)] = root(y);
}

次の図のような２つのグループがあるとき、要素 1 と要素 4 を併合する処理について考えてみます。

グループが上図の左の状態のとき、配列 parent は次のようになっています。

vector<int> parent = {5, 2, -1, -1, 3, 2, 4};

たとえば、要素 0 の親は要素 5 。要素 1 の親は要素 2 です。

unite(1, 4) を実行すれば、どうなるでしょうか。1つずつ見ていきましょう。

unite() 関数のコードを再掲します。

void unite(int x, int y) {
    // x と y が同じグループのときは何もしない
    if (isSame(x, y)) return;

    // x の根の親を y の根にする
    parent[root(x)] = root(y);
}

unite(1, 4) が実行されます。root(1) は 2、 root(4) は 3 です。
３行目の isSame(1, 4) は、要素 1 と要素 4 の根は違うため isSame(1, 4) は false を返します。
そのため、３行目は if (false) になり return は実行されません。
６行目の parent[root(1)] は、root(1) = 2 なので parent[2] になります。
root(4) は 3 なので、parent[2] = 3 になります。

unite(1, 4) を実行したあと、配列 parent は次のようになっています。

 // 併合前
vector<int> parent = {5, 2, -1, -1, 3, 2, 4};

// 併合後
vector<int> parent = {5, 2, 3, -1, 3, 2, 4};

実際にグループが併合できているか確認してみましょう。ためしに、要素 0 の根を確認してみます。

要素 0 の親は、parent[0] の値になるので要素 5
要素 5 の親は、parent[5] の値になるので要素 2
要素 2 の親は、parent[2] の値になるので要素 3
要素 3 の親は、parent[3] の値 (-1) になるので要素 3 が根になる

#include <bits/stdc++.h>
using namespace std;

// 併合前の２つのグループがある状態
vector<int> parent = {5, 2, -1, -1, 3, 2, 4};

int root(int x) {  
    // x が根のとき
    if(parent[x] == -1) return x;

    // 再帰呼び出し（ x の親を引数に ）
    return root(parent[x]);
}

bool isSame(int x, int y) {
    return root(x) == root(y);
}

void unite(int x, int y) {
    // x と y が同じグループのときは何もしない
    if (isSame(x, y)) return;

    // x の根の親を y の根にする
    parent[root(x)] = root(y);
}

int main () {
    cout << "併合前のroot(0)は"<< root(0) << "です。" << endl;

    // 要素 1 と要素 4 の併合
    unite(1, 4);

    cout << "併合後のroot(0)は"<< root(0) << "です。" << endl;
    return 0;
}

// 実行結果
/*
併合前のroot(0)は2です。
併合後のroot(0)は3です。
*/

きちんと併合できているようですね。

簡単な Union-Find の実装コード

説明してきた関数を構造体として実装したのが次のコードです。

struct UnionFind {
    // 親の要素を管理する変数
    vector<int> parent;

    // 変数の初期化
    UnionFind(int n) : parent(n, -1) {}

    // x の根を求める
    int root(int x) {
        // x が根のとき
        if(parent[x] == -1) return x;

        // 再帰呼び出し（ x の親を引数に ）
        return root(parent[x]);
    }

    // x と y の根が同じか
    bool isSame(int x, int y) {
        return root(x) == root(y);
    }

    // x と y のグループを併合する
    void unite(int x, int y) {
        // x と y が同じグループのときは何もしない
        if (isSame(x, y)) return;

        // x の根の親を y の根にする
        parent[root(x)] = root(y);
    }
};

構造体について

初めて出てきた次のコードについて説明します。

// 説明のために一部抜粋
struct UnionFind {

    // 変数の初期化
    UnionFind(int n) : parent(n, -1) {}

};

struct は構造体を表します。

構造体を使えば、複数の型や関数をまとめた新しい型を作ることができます。

たとえば、string 型について考えてみます。string 型の変数 s を宣言するには次のようにします。

string s = "test";

変数 s の文字数は size() を使って取得できます。

cout << s.size()  // => 4

これは、string 型に size() という処理が実装されているため、s.size() で文字数が取得できています。

構造体を使うことで、このような型を自分で作ることができます。

UnionFind という名前の構造体を実装することで、UnionFind 型として扱うことができます。

struct UnionFind {
    /* 省略*/
};

int main () {
    // UnionFind型の変数ufを宣言
    UnionFind uf(n);

    // UnionFind型の isSame() を実行
    uf.isSame(x, y);
}

UnionFind の、初期化処理のコードについて見てみます。次のコードが初期化処理をしているところになります。

struct UnionFind {
    vector<int> parent;

    // 初期化処理をしている部分
    UnionFind(int n) : parent(n, -1) {}

    /* 省略 */
};

構造体の初期化処理は次のように書きます。

struct 構造体名 {
    構造体名(引数) {}
};

初期化処理の中で、構造体の変数を初期化したいときは次のようにします。

struct 構造体名 {
    型　メンバ変数;
    構造体名(引数) : メンバ変数名(初期化内容) {}
};

vector の初期化

UnionFind では、最初の状態ではすべての要素は独立（すべての要素が根）しています。

そのため、要素の親を管理する配列 parent は、すべて -1 （-1は根を表す）である必要があります。

vector は、次のように宣言することで、すべての要素に同じ値を設定できます。

vector<int> v(5, 2);  // => 要素数 5 で、値はすべて 2

コードで確認してみます。

int main () {
    vector<int> v(5, 2);
    
    // v のすべての要素を出力
    for (int i : v) cout << i << " ";
  
    return 0;
}

// 実行結果
/*
2 2 2 2 2 
*/

UnionFind の宣言時に要素数を引数として受け取り、parent (n, -1) ですべての要素を -1 にしています。

struct UnionFind {
    vector<int> parent;

    // 変数の初期化
    UnionFind(int n) : parent(n, -1) {}

    /* 省略 */
};

問題を解いてみる

これで Union-Find のコードの簡単な実装ができました。

AtCoder の AtCoder Typical Contest 001 – B – Union Find で Union-Find を使って問題を解いてみましょう。

先ほど実装した Union-Find を使って、次のように実装しました。

#include<bits/stdc++.h>
using namespace std;

struct UnionFind {
    // 親の要素を管理する変数
    vector<int> parent;

    // 変数の初期化
    UnionFind(int n) : parent(n, -1) {}

    // x の根を求める
    int root(int x) {
        // x が根のとき
        if(parent[x] == -1) return x;

        // 再帰呼び出し（ x の親を引数に ）
        return root(parent[x]);
    }

    // x と y の根が同じか
    bool isSame(int x, int y) {
        return root(x) == root(y);
    }

    // x と y のグループを併合する
    void unite(int x, int y) {
        // x と y が同じグループのときは何もしない
        if (isSame(x, y)) return;

        // x の根の親を y の根にする
        parent[root(x)] = root(y);
    }
};

int main () {
    // 入力
    int N, Q;
    cin >> N >> Q;

    // Union-Find の宣言
    UnionFind uf(N);

    // クエリに答える
    for (int i = 0; i < Q; i++) {

        // 入力
        int P, A, B;
        cin >> P >> A >> B;

        // 連結クエリのとき
        if (P == 0) {
            uf.unite(A, B);
        }

        // 判定クエリのとき
        else {
            if(uf.isSame(A, B)) cout << "Yes" << endl;
            else cout << "No" << endl;
        }
    }

    return 0;
}

このコードで提出したところ、１つのテストケースに TLE してしまいました。

計算量を工夫して実装する

Union-Find の大切な機能は、根を求める root() 関数です。しかし、今の実装だと root() 関数は、根付き木に偏りが発生したときに計算量が多くなってしまいます。

そこで、以下の2つ工夫をして計算量を抑えます。

union by size ( または、union by rank )
経路圧縮

union by size

union by size は、グループを併合するときの処理を工夫します。

次の2つのグループを併合するとき、A グループを根にするか、B グループを根にするかの 2つの方法があります。

要素の数をサイズとしたとき、グループ A はサイズ４。グループ B はサイズ２になります。

併合したあとのそれぞれのグループについて、要素 0 の根を求める root(0) について考えてみます。

A グループを根として併合したとき
1. 要素 0 の親は 5
2. 要素 5 の親は 2
3. 要素 2 は根
B グループを根として併合したとき
1. 要素 0 の親は 5
2. 要素 5 の親は 2
3. 要素 2 の親は 3
4. 要素 3 は根

A グループを根として併合したときの方が、根を求めるための計算回数が少なくなりました。

このように、２つのグループを併合しただけですが、どちらのグループを根にするかによって、根を求める計算回数に違いがでました。

実は、サイズの小さいグループの根の親を大きいグループの根に併合することで、根を求める計算量を少なくできます。

根を求める処理は、Union-Find において重要な処理です。根を求める計算量は少なければ少ないほど効率的になります。

「サイズ（要素数）の小さいの方の根付き木を、大きい方に併合する」という考え方を union by size といいます。

union by size の実装

union by size は、サイズの違いでどちらを根にするかを判断します。

そのため、グループのサイズを管理するための配列 size を導入します。

始めは、各要素はバラバラの単独グループなので、サイズは１です。

struct UnionFind {
    // 親の要素とサイズを管理する変数
    vector<int> parent, size;

    // 変数の初期化 ( size のすべての値を 1 にしている )
    UnionFind(int n) : parent(n, -1), size(n, 1) {}

    // 省略
};

併合するときに、サイズが小さい方の親を、サイズの大きい方の根にします。

unite() 関数の変更前（左）と変更後（右）のコードです。

変更後のコードでは、unite() 関数の最初に x と y の根を変数に代入して、根を何度も求めないようにもしています。

// 変更前
void unite(int x, int y) {
    // x と y が同じグループのときは何もしない
    if (isSame(x, y)) return;

    // x の根の親を y の根にする
    parent[root(x)] = root(y);
}

// 変更後
// x と y のグループを併合する
void unite(int x, int y) {
    // x と y の根を取得
    int rootX = root(x);
    int rootY = root(y);

    // x と y が同じグループのときは何もしない
    if (rootX == rootY) return;

    // union by size
    // x のサイズが小さいとき
    if (size[rootX] < size[rootY]) {
        // x の親を y にする
        parent[rootX] = rootY;

        // y のサイズに x のサイズを足す
        size[rootY] += size[rootX]; 
    }
    // y のサイズが小さいとき
    else {
        // y の親を x にする
        parent[rootY] = rootX;

        // x のサイズに y のサイズを足す
        size[rootX] += size[rootY]; 
    }
}

union by size の部分は swap() を使えば次のように簡潔に書けます。

 // x と y のグループを併合する
    void unite(int x, int y) {
        // x と y の根を取得
        int rootX = root(x);
        int rootY = root(y);
      
        // x と y が同じグループのときは何もしない
        if (rootX == rootY) return;

        // union by size（ y のサイズが小さくなるように調整 ）
        if (size[rootX] < size[rootY]) swap(rootX, rootY);
      
        // y の親を x にする
        parent[rootY] = rootX;
      
        // x のサイズに y のサイズを足す
        size[rootX] += size[rootY]; 
    }

swap(a, b) は、a と b の値を入れ替えます。

#include<bits/stdc++.h>
using namespace std;

int main () {
    int a = 1;
    int b = 3;

    // a と b の値を入れ替える
    swap(a, b);

    cout << a << endl;  // => 3
    cout << b << endl;  // => 1

    return 0;
}

経路圧縮

経路圧縮は、根を求めるときの root() 関数についての工夫です。

次の図では、要素 1 から要素 6 までの根は要素 0 で、配列 parent は次のようになります。

vector<int> parent = {-1, 2, 0, 1, 1, 3, 3};  // -1 は根を表す

root(1) から root(6) の戻り値は、すべて 0 になります。

#include<bits/stdc++.h>
using namespace std;

// 最小限のコードのみの抜粋
vector<int> parent = {-1, 2, 0, 1, 1, 3, 3};

int root(int x) {
    // x が根のとき
    if(parent[x] == -1) return x;
    
    // 再帰呼び出し（ x の親を引数に ）
    return root(parent[x]);
}

int main () {
    cout << root(1) << endl;  // => 0
    cout << root(2) << endl;  // => 0
    cout << root(3) << endl;  // => 0
    cout << root(4) << endl;  // => 0
    cout << root(5) << endl;  // => 0
    cout << root(6) << endl;  // => 0
    return 0;
}

経路圧縮は、root(x) を実行したときに、要素 x から根までの要素の親を根に置き換えます。

先ほどの図で root(6) を実行したとき、根に行くまでに通った要素 3, 1, 2 の親を根である要素 0 に置き換えます。

root(6) を実行して経路圧縮した結果、配列 parent は次のようになります。

// root(6) 実行前
vector<int> parent = {-1, 2, 0, 1, 1, 3, 3};

// root(6) 実行後
vector<int> parent = {-1, 0, 0, 0, 1, 3, 0};

要素 4 と要素 5 は、経路として通っていないため、親は変わりません。

経路圧縮の実装

経路圧縮は root() 関数を変更することで実現します。

// 変更前
int root(int x) {
    // x が根のとき
    if(parent[x] == -1) return x;

    // 再帰呼び出し（ x の親を引数に ）
    return root(parent[x]);
}

// 変更後
int root(int x) {
    // x が根のとき
    if(parent[x] == -1) return x;

    // 経路圧縮
    return parent[x] = root(parent[x]);
}

７行目のコードが変わっています。このコードで、要素 x の親要素を根に変更しています。

変更後のコードがわかりにくければ、次のコードを見てください。６行目で x の根を求めています。　

int root(int x) {
    // x が根のとき
    if(parent[x] == -1) return x;
    
    // 再帰呼び出しで根を取得
    int rootX = root(parent[x]);

    // x の親を根に変更
    parent[x] = rootX;
     
    // 根を返す
    return rootX;
}

これで root() 関数が呼び出されたとき、要素 x から根までの経路の親が根に置き換えられます。

この記事での最終の実装コード

以上をふまえて、この記事の一番初めに紹介した Union-Find のコードが実装できました。

struct UnionFind {
    // 親の要素とサイズを管理する変数
    vector<int> parent, size;

    // 変数の初期化
    UnionFind(int n) : parent(n, -1), size(n, 1) {}

    // x の根を求める
    int root(int x) {
        // x が根のとき
        if(parent[x] == -1) return x;

        // 経路圧縮
        return parent[x] = root(parent[x]);
    }

    // x と y の根が同じか
    bool isSame(int x, int y) {
        return root(x) == root(y);
    }

    // x と y のグループを併合する
    void unite(int x, int y) {
        // x と y の根を取得
        int rootX = root(x);
        int rootY = root(y);

        // x と y が同じグループのときは何もしない
        if (rootX == rootY) return;

        // union by size（ y のサイズが小さくなるように調整 ）
        if (size[rootX] < size[rootY]) swap(rootX, rootY);

        // y の親を x にする
        parent[rootY] = rootX;

        // x のサイズに y のサイズを足す
        size[rootX] += size[rootY]; 
    }
};

このコードで AtCoder Typical Contest 001 – B – Union Find を再び解いてみます。

#include<bits/stdc++.h>
using namespace std;

struct UnionFind {
    // 親の要素とサイズを管理する変数
    vector<int> parent, size;

    // 変数の初期化
    UnionFind(int n) : parent(n, -1), size(n, 1) {}

    // x の根を求める
    int root(int x) {
        // x が根のとき
        if(parent[x] == -1) return x;

        // 経路圧縮
        return parent[x] = root(parent[x]);
    }

    // x と y の根が同じか
    bool isSame(int x, int y) {
        return root(x) == root(y);
    }

    // x と y のグループを併合する
    void unite(int x, int y) {
        // x と y の根を取得
        int rootX = root(x);
        int rootY = root(y);

        // x と y が同じグループのときは何もしない
        if (rootX == rootY) return;

        // union by size（ y のサイズが小さくなるように調整 ）
        if (size[rootX] < size[rootY]) swap(rootX, rootY);

        // y の親を x にする
        parent[rootY] = rootX;

        // x のサイズに y のサイズを足す
        size[rootX] += size[rootY]; 
    }
};

int main () {
    // 入力
    int N, Q;
    cin >> N >> Q;

    // Union-Find の宣言
    UnionFind uf(N);

    // クエリに答える
    for (int i = 0; i < Q; i++) {

        // 入力
        int P, A, B;
        cin >> P >> A >> B;

        // 連結クエリのとき
        if (P == 0) {
            uf.unite(A, B);
        }

        // 判定クエリのとき
        else {
            if(uf.isSame(A, B)) cout << "Yes" << endl;
            else cout << "No" << endl;
        }
    }

    return 0;
}

このコードで提出したところ、無事 AC になりした。

まとめ

Union-Find はグループ分けを効率的に管理できるデータ構造です。
次の2つのことが効率的に行えます。
- 要素 x と要素 y のグループを併合する
- 2つの要素 x, y が同じグループに属するか調べる
Union-Find は、根付き木を使って実現する。
2つの要素が同じグループかは、それぞれの要素の根が同じかどうかで判定できる。
2つのグループを併合するときは、片方の根をもう1つのグループの根の子になるようにつなぐ
union by size と経路圧縮をすることで、計算量を抑えることができる。
union by size は、併合するときの工夫。サイズ（頂点の数）が小さい方の根を、サイズの大きい方の根につなげる。
経路圧縮は、根を求めるときの工夫。要素 x の根を求めるときに、その経路の要素を根に張りかえる。

参考

以下の書籍を参考にさせていただきました。

この中でも、問題解決力を鍛える!アルゴリズムとデータ構造がUnion-Findについてわかりやすく説明されていると思います。

リンク

次の記事もおすすめです。