7K4B blog

猫でも分かる何か

NyaaLIB::DS_UnionFind

Union Find

  • DS_UnionFind(n): 半開区間[0,n) のデータ構造を確保

 O(\log N) 以下で処理できる関数

  • Union(x, y): xとyを結合
  • Find(x): xの根を返す
  • Max(x): xが属する集合の最大値を取得
  • Min(x): xが属する集合の最小値を取得
  • Same(x, y): xとyが同じ集合どうか判定
  • Size(x): xが属する集合サイズ

 O(N \log N) で処理する関数

  • VertexSet(l, r): [l, r)区間における頂点集合を取得

f:id:tkr987:20220302211937p:plain
SameとUnionを組み合わせると閉路検出になる(有向無向どちらでも使える)

#include <bits/stdc++.h>

namespace NyaaLIB
{
	/**
	* UnionFind ライブラリ
	* O(log N) 以下でマージできるデータ構造
	**/

	struct UnionFindVertex
	{	// 頂点を表現する構造体
		using ll = long long;
		ll self = 0; // 自分自身のインデックス
		ll root = 0; // 根のインデックス
		ll size = 0; // 自分が属している木のサイズ
	};

	struct DS_UnionFind
	{	// 木構造で素集合データ構造を実現させる
		using ll = long long;
		std::vector<UnionFindVertex> v;
		std::vector<ll> smin; // 集合の最小値
		std::vector<ll> smax; // 集合の最大値

		DS_UnionFind(ll n) : v(n), smin(n), smax(n)
		{	// 半開区間[0,n)の素集合データ構造を作成する
			for (ll i = 0; i < n; i++)
			{	// 各頂点の根を自分自身で初期化、木サイズは1
				v[i].self = i;
				v[i].root = i;
				v[i].size = 1;
				smin[i] = smax[i] = i;
			}
		}

		ll Find(ll x)
		{	// 要素xの根を検索する
			if (x == v[x].root) return x;
			// 根の探索をすると同時に次からO(1)で根を参照できるようにする(経路圧縮)
			v[x].root = Find(v[x].root);
			return v[x].root;
		}

		ll Min(ll x) { return smin[Find(x)]; }
		ll Max(ll x) { return smax[Find(x)]; }
		bool Same(ll x, ll y) { return Find(x) == Find(y); }
		ll Size(ll x) { return v[Find(x)].size; }

		/**
		* @brief 要素を併合する関数
		* @note
		* 要素xを含む木と要素yを含む木を「サイズ優先で」併合する。
		* ただし、xとyが既に同じ木に属しているときは何もしない。
		* 併合したときtrue、何もしなかったときfalseを返す。
		* サイズによる工夫により、計算量はアッカーマンの逆関数になる。
		**/
		bool Union(ll x, ll y)
		{	// 既に同じ木に属しているときは何もしない
			ll root1 = Find(x);
			ll root2 = Find(y);
			if (root1 == root2) return false;

			UpdateMin(root1, root2);
			UpdateMax(root1, root2);
			// サイズの小さい木の根をサイズの大きい木の根に繋いで併合する
			if (v[root1].size < v[root2].size)
			{
				v[root1].root = root2;
				v[root2].size += v[root1].size;
			}
			else
			{
				v[root2].root = root1;
				v[root1].size += v[root2].size;
			}
			return true;
		}

		/**
		* @brief 要素を併合する関数
		* @param x 併合する要素1
		* @param y 併合する要素2
		* @param p  親指定
		* @note
		* 要素xを含む木と要素yを含む木を「pを含む木を親として」併合する。
		* ただし、xとyが既に同じ木に属しているときは何もしない。
		* 併合したときtrue、何もしなかったときfalseを返す。
		* 計算量はO(logN)になり、アッカーマンの逆関数に比べて若干遅くなる。
		**/
		bool Union(ll x, ll y, ll p)
		{	// 既に同じ木に属しているときは何もしない
			ll root1 = Find(x);
			ll root2 = Find(y);
			ll rootp = Find(p);
			if (root1 == rootp && root2 == rootp) return false;

			UpdateMin(root1, root2);
			UpdateMax(root1, root2);
			// 子の木を親の木へ併合する
			if (rootp == root1)
			{
				v[root2].root = rootp;
				v[rootp].size += v[root2].size;
			}
			else if (rootp == root2)
			{
				v[root1].root = rootp;
				v[rootp].size += v[root1].size;
			}
			else
			{
				v[root1].root = rootp;
				v[rootp].size += v[root1].size;
				v[root2].root = rootp;
				v[rootp].size += v[root2].size;
			}
			return true;
		}

		std::set<ll> VertexSet(ll l, ll r)
		{	// 半開区間[l,r)における頂点集合を返す、計算量 O(N log N)
			std::set<ll> res;
			for (ll i = l; i < r; i++) res.insert(Find(i));
			return res;
		}
	private:
		void UpdateMin(ll x, ll y) { smin[x] = smin[y] = std::min(smin[x], smin[y]); }
		void UpdateMax(ll x, ll y) { smax[x] = smax[y] = std::max(smax[x], smax[y]); }
	};
}

実装例
https://atcoder.jp/contests/abc231/submissions/30370697
https://atcoder.jp/contests/abc226/submissions/30370569