Union-Findの実装

Union-Findのアルゴリズムを勉強したので、競技プログラミングなどで必要になった時に、すぐに使えるようにメモを残しておきたいと思います。

概要

Union-Findとは、グループ分けを効率的に管理することができるデータ構造です。Union-Findを使用することによって、異なるグループを統合したり、2つの要素が同じグループに含まれているかを判定することができます。実現したいことのイメージは以下の図のようになります。

https://algo-method.com/descriptions/132

実装

Union-FindのC++での実装例を載せておきます。

UnionFind.cpp

#include <iostream>
#include <array>

class UnionFind
{
public:
	std::array<int, 100009> par; //各要素の親ノードを管理する配列
	std::array<int, 100009> siz; //頂点xを根とするグループの要素数を管理する配列

	UnionFind()
	{
		par.fill(-1); //最初は親ノードを持たないので-1で初期化
		siz.fill(1); //最初は各要素がそれぞれ異なるグループに属しているので1で初期化
	}

    //頂点xの根を求める
	int root(int x)
	{
		while (par.at(x) != -1)
		{
			x = par.at(x); //親ノードがなくなるまで探索する
		}

		return x;
	}

    //要素uとvのグループを統合する
	void unite(int u, int v)
	{
		int RootU = root(u);
		int RootV = root(v);

		if (RootU == RootV)
		{
			return;
		}

        //要素数が多いグループの根を統合後のグループの根にする
		if (siz.at(RootU) < siz.at(RootV))
		{
			par.at(RootU) = RootV;
			siz.at(RootV) = siz.at(RootU) + siz.at(RootV);
		}
		else
		{
			par.at(RootV) = RootU;
			siz.at(RootU) = siz.at(RootU) + siz.at(RootV);
		}
	}

    //要素uとvが同じグループに属するかを判定する
	bool same(int u, int v)
	{
		if (root(u) == root(v))
		{
			return true;
		}
		else
		{
			return false;
		}
	}

    //要素xのグループのメンバーの数を求める
	int count_member(int x)
	{
		int RootX = root(x);
		return siz.at(RootX);
	}
};

今回の記事はメモの代わりなので、Union-Findの詳しい仕組みを理解したい方は書籍などを読むことをお勧めします。

使い方

Union-Findクラスの使い方を以下のコードのmain関数の中に示しています。

UnionFind.cpp

#include <iostream>
#include <array>

class UnionFind
{
public:
	std::array<int, 100009> par;
	std::array<int, 100009> siz;

	UnionFind()
	{
		par.fill(-1);
		siz.fill(1);
	}

	int root(int x)
	{
		while (par.at(x) != -1)
		{
			x = par.at(x);
		}

		return x;
	}

	void unite(int u, int v)
	{
		int RootU = root(u);
		int RootV = root(v);

		if (RootU == RootV)
		{
			return;
		}

		if (siz.at(RootU) < siz.at(RootV))
		{
			par.at(RootU) = RootV;
			siz.at(RootV) = siz.at(RootU) + siz.at(RootV);
		}
		else
		{
			par.at(RootV) = RootU;
			siz.at(RootU) = siz.at(RootU) + siz.at(RootV);
		}
	}

	bool same(int u, int v)
	{
		if (root(u) == root(v))
		{
			return true;
		}
		else
		{
			return false;
		}
	}

	int count_member(int x)
	{
		int RootX = root(x);
		return siz.at(RootX);
	}
};

int main()
{
	UnionFind uf;

	std::cout << std::boolalpha << uf.same(1, 2) << std::endl;
	std::cout << uf.count_member(2) << std::endl;
	uf.unite(1, 2);
	std::cout << uf.same(1, 2) << std::endl;
	std::cout << uf.count_member(2) << std::endl;
}

出力

false
1
true
2

unite関数で1と2を結合した結果が反映されていることがわかります。
Union-Findはどの拠点と拠点が結ばれているかとか、属しているコミュニティのフレンドの数とか管理するのに使えそうかと思います。

参考