そまちょブログのそまちょ(@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 と経路圧縮を実装して上のコードになります。
基本的な考え方
7人でグループ分けをすることになりました。各グループには、1人リーダーがいます。
グループ分けをした結果、次の3つのグループができました。それぞれのグループのリーダーは ②、③、④ です。
①のリーダーは②で、⑤のリーダーも②です。
⓪のリーダーは②ですが、⑥のリーダーは④です。
リーダーが同じとき、その人たちは同じグループ。リーダーが違うとき、その人たちは違うグループということになります。
次に、③がリーダーのグループと⑥の所属するグループが併合しました。これにより、④と⑥のリーダーは③に変わりました。
併合するときは、併合するグループの各要素について、新しいリーダーを設定すればいいということになります。
- リーダーが同じなら同じグループ
- リーダーが別なら違うグループ
- グループが併合したら、どちらかのグループのメンバーのリーダが変わる
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 で管理します。こうすることで、ある要素から順々に親をたどっていけば、いつかは根にたどり着きます。
次の図は、2つの根付き木のグループに分かれているときの、0から6までの各要素の配列 parent を表しています。
これをふまえて、要素 x の根を求める root() 関数のコードは次のようになります。配列 parent で、根は -1 で表します。
int root(int x) {
// x が根のとき
if(parent[x] == -1) return x;
// 再帰呼び出し( x の親を引数に )
return root(parent[x]);
}
このコードでは、根を再帰的に求めています( 6行目の return root(parent[x]); のところ )。
たとえば、配列 parent が次の値だったとき、要素 0 の根を求めるために root(0) を呼び出すとどうなるでしょうか。
vector<int> parent = {5, 2, -1, -1, 3, 2, 4};
- root(0) が呼び出される。parent[0] は 5 のため6行目は return root(5) になる。
- root(5) が呼び出される。parent[5] は 2 のため6行目は return root(2) になる。
- root(2) が呼び出される。parent[2] は -1 のため3行目の 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 の根が求めれました。
2つの要素の根が同じかどうかを判定すれば、グループが同じかわかります。
同じグループかを判定する 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つの要素のグループは別グループである必要があります。
そのため、2つの要素の根が同じときは、すでに同じグループなので併合する必要はありません。
グループを併合するときは、片方のグループの根の親を、もう片方のグループの根になるように設定します。
これをふまえて、実装した unite() 関数が次のコードです。
void unite(int x, int y) {
// x と y が同じグループのときは何もしない
if (isSame(x, y)) return;
// x の根の親を y の根にする
parent[root(x)] = root(y);
}
次の図のような2つのグループがあるとき、要素 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 です。
- 3行目の isSame(1, 4) は、要素 1 と要素 4 の根は違うため isSame(1, 4) は false を返します。
- そのため、3行目は if (false) になり return は実行されません。
- 6行目の 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;
// 併合前の2つのグループがある状態
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;
}
このコードで提出したところ、1つのテストケースに TLE してしまいました。
計算量を工夫して実装する
Union-Find の大切な機能は、根を求める root() 関数です。しかし、今の実装だと root() 関数は、根付き木に偏りが発生したときに計算量が多くなってしまいます。
そこで、以下の2つ工夫をして計算量を抑えます。
- union by size ( または、union by rank )
- 経路圧縮
union by size
union by size は、グループを併合するときの処理を工夫します。
次の2つのグループを併合するとき、A グループを根にするか、B グループを根にするかの 2つの方法があります。
要素の数をサイズとしたとき、グループ A はサイズ4。グループ B はサイズ2になります。
併合したあとのそれぞれのグループについて、要素 0 の根を求める root(0) について考えてみます。
- A グループを根として併合したとき
- 要素 0 の親は 5
- 要素 5 の親は 2
- 要素 2 は根
- B グループを根として併合したとき
- 要素 0 の親は 5
- 要素 5 の親は 2
- 要素 2 の親は 3
- 要素 3 は根
A グループを根として併合したときの方が、根を求めるための計算回数が少なくなりました。
このように、2つのグループを併合しただけですが、どちらのグループを根にするかによって、根を求める計算回数に違いがでました。
実は、サイズの小さいグループの根の親を大きいグループの根に併合することで、根を求める計算量を少なくできます。
根を求める処理は、Union-Find において重要な処理です。根を求める計算量は少なければ少ないほど効率的になります。
「サイズ(要素数)の小さいの方の根付き木を、大きい方に併合する」という考え方を union by size といいます。
union by size の実装
union by size は、サイズの違いでどちらを根にするかを判断します。
そのため、グループのサイズを管理するための配列 size を導入します。
始めは、各要素はバラバラの単独グループなので、サイズは1です。
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]);
}
7行目のコードが変わっています。このコードで、要素 x の親要素を根に変更しています。
変更後のコードがわかりにくければ、次のコードを見てください。6行目で 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についてわかりやすく説明されていると思います。
次の記事もおすすめです。