[競プロ]Union-Find実装【Java】

競技プログラミングで頻出のUnion-Findというデータ構造について学びました。
また、例題を解きましたので解説していきます。

Union-Findって？

データ型の一種です。
wikipediaを見ると以下のように書かれています。

素集合データ構造（そしゅうごうデータこうぞう、英: disjoint-set data structure）は、データの集合を素集合（互いにオーバーラップしない集合）に分割して保持するデータ構造。このデータ構造に対する以下の2つの便利な操作をUnion-Findアルゴリズムと呼ぶ。

Find: 特定の要素がどの集合に属しているかを求める。2つの要素が同じ集合に属しているかの判定にも使われる。
Union: 2つの集合を1つに統合する。
これら2つの操作をサポートしているため、素集合データ構造は「Union-Findデータ構造」あるいは「Merge-Find集合」とも呼ばれる。

「任意の2要素が同じ集合に属している」ことを判定したり、
「2つの集合を合体する」ことが出来たりするようです。Union, Find合わせてUnionFindというみたいですね。

集合の中の要素、それぞれに対して親を定義して、2つの要素で根まで辿ってみてそれが同じだったら2つの要素は同じ集合に属している、また、2つの集合をあわせるときは一方の根の親をもう一方の根に紐づけてあげると良い、ということです。

正直、詳しい説明は出来ないので、参考文献を見るとわかりやすい説明が載ってます。
この記事ではJavaで、UnionFindを簡単に書いてみて、そのコードを紹介していこうと思います。

Javaで実装

以下のとおりです。主にAtCoder公式のスライド（参考文献）を参考に作成しました。
スライドで触れられている経路圧縮やrankの考慮も入れています。

class UnionFind {
	int[] parent;
	int[] rank;

	public UnionFind(int n) {
		// 初期化コンストラクタ
		this.parent = new int[n];
		this.rank = new int[n];

		// 最初はすべてが根
		for (int i = 0; i < n; i++) {
			parent[i] = i;
			rank[i] = 0;
		}
	}

	/**
	 * 要素の根を返す。
	 * 経路圧縮付き。（1→3→2となっていて2をfindした際、1→3,2と木の深さを浅くする。）
	 *
	 * @param x
	 * @return 要素xの根
	 */
	public int find(int x) {
		if (x == parent[x]) {
			return x;
		} else {
			// 経路圧縮時はrank変更しない
			parent[x] = find(parent[x]);
			return parent[x];
		}
	}

	/**
	 * ２つの要素が同じ集合に属するかどうかを返す。
	 *
	 * @param x
	 * @param y
	 * @return 同じ集合であればtrue
	 */
	public boolean same(int x, int y) {
		return find(x) == find(y);
	}

	/**
	 * 要素xが属する集合と要素yが属する集合を連結する。
	 * 木の高さ（ランク）を気にして、低い方に高い方をつなげる。（高い方の根を全体の根とする。）
	 *
	 * @param x
	 * @param y
	 */
	public void unite(int x, int y) {
		int xRoot = find(x);
		int yRoot = find(y);

		if (xRoot == yRoot) {
			// 属する集合が同じな場合、何もしない
			return;
		}

		// rankを比較して共通の根を決定する。
		// ※find時の経路圧縮はrank考慮しない
		if (rank[xRoot] > rank[yRoot]) {
			// xRootのrankのほうが大きければ、共通の根をxRootにする
			parent[yRoot] = xRoot;
		} else if (rank[xRoot] < rank[yRoot]) {
			// yRootのrankのほうが大きければ、共通の根をyRootにする
			parent[xRoot] = yRoot;
		} else {
			// rankが同じであれば、どちらかを根として、rankを一つ上げる。
			parent[xRoot] = yRoot;
			rank[xRoot]++;
		}
	}
}

問題演習

丁度いいのでこちらで演習してみましょう。
atc001 B - Union Find

the・UnionFindの問題です。以下のコードでACでした。

import java.util.Scanner;

public class Main {

	public static void main(String[] args) {
		Scanner sc = new Scanner(System.in);

		int n = sc.nextInt();
		int q = sc.nextInt();

		UnionFind uf = new UnionFind(n);

		for (int i = 0; i < q; i++) {
			if (sc.nextInt() == 0) {
				// p = 0の場合、要素を接続
				uf.unite(sc.nextInt(), sc.nextInt());
			} else {
				// p = 1の場合、要素が接続されているかを出力
				if (uf.same(sc.nextInt(), sc.nextInt())) {
					System.out.println("Yes");
				} else {
					System.out.println("No");
				}
			}
		}
	}
}

class UnionFind {
	int[] parent;
	int[] rank;

	public UnionFind(int n) {
		// 初期化コンストラクタ
		this.parent = new int[n];
		this.rank = new int[n];

		// 最初はすべてが根
		for (int i = 0; i < n; i++) {
			parent[i] = i;
			rank[i] = 0;
		}
	}

	/**
	 * 要素の根を返す。
	 * 経路圧縮付き。（1→3→2となっていて2をfindした際、1→3,2と木の深さを浅くする。）
	 *
	 * @param x
	 * @return 要素xの根
	 */
	public int find(int x) {
		if (x == parent[x]) {
			return x;
		} else {
			// 経路圧縮時はrank変更しない
			parent[x] = find(parent[x]);
			return parent[x];
		}
	}

	/**
	 * ２つの要素が同じ集合に属するかどうかを返す。
	 *
	 * @param x
	 * @param y
	 * @return 同じ集合であればtrue
	 */
	public boolean same(int x, int y) {
		return find(x) == find(y);
	}

	/**
	 * 要素xが属する集合と要素yが属する集合を連結する。
	 * 木の高さ（ランク）を気にして、低い方に高い方をつなげる。（高い方の根を全体の根とする。）
	 *
	 * @param x
	 * @param y
	 */
	public void unite(int x, int y) {
		int xRoot = find(x);
		int yRoot = find(y);

		if (xRoot == yRoot) {
			// 属する集合が同じな場合、何もしない
			return;
		}

		// rankを比較して共通の根を決定する。
		// ※find時の経路圧縮はrank考慮しない
		if (rank[xRoot] > rank[yRoot]) {
			// xRootのrankのほうが大きければ、共通の根をxRootにする
			parent[yRoot] = xRoot;
		} else if (rank[xRoot] < rank[yRoot]) {
			// yRootのrankのほうが大きければ、共通の根をyRootにする
			parent[xRoot] = yRoot;
		} else {
			// rankが同じであれば、どちらかを根として、rankを一つ上げる。
			parent[xRoot] = yRoot;
			rank[xRoot]++;
		}
	}
}

//

参考文献