【アルゴリズム紹介】Union-Find木
Page content
背景
3月上旬ごろに開かれたAtCoderの「AtCoder Beginner Contest 157」にて、D問題が解けず悔しさを抱きつつ解法を調べました。
その際Union-Find木という新たなデータ構造を知り、今後も使いそうだったため知識の定着も兼ねてこちらでまとめることにしました。
コンテスト中に解けなかった問題はこちらです。
ABC157 D - Friend Suggestion
用途
Union-Find木は、要素のリンクをもとに、リンクしている要素の集合を生成します。
以下の図のように、リンクで繋がっている要素をまとめた集合を生成します。
Union-Find木で生成された集合を用い、2要素が同じ集合にいるか調べることができます。
また、一部改修することで各集合の要素数を求めることもできます。
コード
#include <bits/stdc++.h>
using namespace std;
class UnionFind {
public:
UnionFind(int N) : root(N), tree_rank(N, 0), tree_size(N, 1) {
for (int i = 0; i < N; i++) root[i] = i;
}
void MergeTree(int x, int y) {
x = FindRoot(x);
y = FindRoot(y);
if (x == y) return;
if (tree_rank[x] < tree_rank[y]) swap(x, y);
root[y] = x;
tree_size[x] += tree_size[y];
if (tree_rank[x] == tree_rank[y]) tree_rank[x]++;
}
bool IsSameTree(int x, int y) { return (FindRoot(x) == FindRoot(y)); }
int GetTreeSize(int x) { return tree_size[FindRoot(x)]; }
private:
vector<int> root;
vector<int> tree_rank;
vector<int> tree_size;
int FindRoot(int x) {
if (root[x] == x) return x;
root[x] = FindRoot(root[x]);
return root[x];
}
};
使い方
まずツリーの宣言とリンクの設定を行います。
宣言とリンクの設定は以下のようなコードになります。
// ツリーの宣言
int N = 7;
UnionFind tree(N);
// リンク設定
tree.MergeTree(0, 1);
tree.MergeTree(1, 2);
tree.MergeTree(3, 4);
tree.MergeTree(4, 5);
tree.MergeTree(5, 6);
説明
参考元のスライドが一番分かりやすかったため、そちらを紹介します。