UnionFind木の紹介

導入

図のように島があります。図の灰色の線分のところに橋を作ることができますが、作るにはその線分に書かれた数字だけコストがかかります。どの島からどの島にも行き来できるように作るべき橋を選ぶ方法のうち、コストが最小のものを求めてください。

この問題はコストが小さい橋から順番に、その橋が必要であれば選ぶという操作を繰り返せば解けます（クラスカル法と呼ばれるアルゴリズムです）。橋が必要かどうかは、その時点でその橋の両端が行き来不可能かどうかということですが、これはどのように判定すればいいでしょうか。そこで出てくるのがUnionFind木というデータ構造です。

UnionFind木

UnionFind木はUnion, Findという操作を効率的に行うためのデータ構造です。

$n$個の集合$\lbrace 1\rbrace, \lbrace 2\rbrace, \ldots , \lbrace n\rbrace$に対して、以下の操作を行います。

• Union(x, y): xが含まれる集合とyが含まれる集合をマージする
• Find(x): xが含まれる集合を取得する

UnionFind木では、それぞれの集合を木構造で保持します。木の根を集合の代表として、

• union(x,y): xが属する木とyが属する木をマージする
• find(x): xが属する木の根の番号を取得する
  としています。

Union(x, y): xが属する木とyが属する木をマージする

配列parentを用いて、iの親の番号をparent[i]で保持します。すると、木のマージは根ノードの親をもう一方の木のノードにすることで実現できます。

self.parent[self.find(x)] = y;  // yが属する木にxが属する木をマージする

find(x)のことを考えると木の高さがなるべく小さくなるようにマージする必要があります。なので、yではなく、yの根ノードにくっつけた方がいいです。

self.parent[self.find(x)] = self.find(y);  // yの木の根ノードにxの木をぶら下げる

さらに、木の高さが高い方に低い方をぶら下げた方が、結果の木の高さを小さくすることができます。
なので高さを管理する配列rankを用意して、以下のようにマージするのがいいです。（この工夫はunion by rankと呼ばれます。かわりに木の高さではなく大きさで判断するunion by sizeという方法も存在するようです。）

let root_x = self.find(x);
let root_y = self.find(y);

let rank_x = self.rank[root_x];
let rank_y = self.rank[root_y];

if rank_x <= rank_y {
  // yにxをマージする
  self.parent[root_x] = root_y;

  // 高さが等しいとき、マージ後のyの木の高さは+1される
  if rank_x == rank_y {
    self.rank[root_y] += 1;
  }
} else {
  // xにyをマージする
  self.parent[root_y] = root_x;

  // rank_x > rank_yなので、xの木の高さは変わらない
}

find(x): xが属する木の根の番号を取得する

根ノードはparent[i] == iのようになっているので、これが見つかるまで親を探し続けます。

fn find(&self, x: usize) {
    if self.parent[x] == x {
        x
    } else {
        self.find(self.parent[x])
    }
}

このデータ構造では根のノードにしか興味がなく、親はぶっちゃけどうでもいいので、parent[i]を更新してしまえば次回以降は高速にアクセスできます。（この工夫はpath compressionと呼ばれます。)

fn find(&mut self, x: usize) {
    if self.parent[x] == x {
        x
    } else {
        let root = self.find(self.parent[x]);
        self.parent[x] = root;
        root
    }
}

実装例

struct UnionFind {
	parent: Vec<usize>,
	rank: Vec<usize>
}

impl UnionFind {
    // 初期化: {0}, {1}, ..., {n-1} を作る
	pub fn new(n: usize) -> Self {
		Self {
			parent: (0..n).collect(),
			rank: vec![0; n]
		}
	}

    pub fn find(&mut self, x: usize) {
        if self.parent[x] == x {
            x
        } else {
            let root = self.find(self.parent[x]);
            self.parent[x] = root;
            root
        }
    }

    // xとyが同じ集合に属しているかどうかを判定する
    pub fn same(&mut self, x: usize, y: usize) {
        self.find(x) == self.find(y)
    }


    pub fn union(&mut self, x: usize, y: usize) {
        if self.same(x, y) {
            return;
        }
        let root_x = self.find(x);
        let root_y = self.find(y);

        let rank_x = self.rank[root_x];
        let rank_y = self.rank[root_y];

        if rank_x <= rank_y {
            // yにxをマージする
            self.parent[root_x] = root_y;

            // 高さが等しいとき、マージ後のyの木の高さは+1される
            if rank_x == rank_y {
                self.rank[root_y] += 1;
            }
        } else {
            // xにyをマージする
            self.parent[root_y] = root_x;

            // rank_x > rank_yなので、xの木の高さは変わらない
        }
    }
}

ここからさらに、要素数や最大値など集合の統計情報を持たせてもいいかもしれません。（自分は必要な時はUnionFindとは別に統計情報のための配列を管理していますが、unionでどちらの木が代表になるかが利用側はわからないので、実装しておくと便利だなぁと思っています。）