Zenn
📖

グルーピングのためのデータ構造：Union Find Tree

2024/05/19に公開
Rust
競技プログラミング
アルゴリズムとデータ構造
tech

所属するグループがわかっている場合

単に２つのグループに分ける場合、標準ライブラリにあるpartition()関数が便利です。また、「血液型がO型の人」など特定の要素を取り出すためにはfilter()関数やその亜種を用いると良いです。３つ以上のグループに分割する際は、ループとmatch式を組み合わせると良いです。

２つの要素の関係が与えられる場合

n人の生徒からなるクラスがある。m個の友達関係に基づいてお友達グループをつくれ。
友達関係は２名の生徒が友達であることをを表す。また、友達の友達は友達であるとする。

こうした２要素間の関係にもとづいてグルーピングする場合も、for文と条件式を組み合わせでなんとかできます。問題文をアルゴリズムっぽく言い換えてみましょう。

i番目の友達関係(A_i, B_i)が与えられたとする。A_iさんとB_iさんが同じお友達グループに所属していれば何もしない。別のグループに所属していれば２つのグループを結合する。これをn回繰り返す。

ここで登場する「同じグループに所属しているか判定する」「２つのグループを結合する」などの操作を高速[1]に実現するのが Union Find Tree と呼ばれるデータ構造です。

データ構造

Union Find Tree ではグループを木として管理します。子ノードは親ノードのインデックスをもち、ルートは木のサイズをもちます。Union Find Tree 全体で森を構成します。

お友達グループのリーダーがグループの人数をもちます。それ以外のメンバーはリーダーか自身以外のメンバーのインデックスをもちます。

「木」とは

n個のノード(頂点)とn-1個の辺からなる連結なグラフ。連結とはどの２つのノードの間にも道があることをいう。道は連続する辺の集まりのこと。図の矢印は親から子に向かって伸びている。aは親をもたない特別なノードでルート(根)という。

#[derive(Debug, Clone, Copy)]
enum Node {
    Size(usize),
    Parent(usize),
}

#[derive(Debug)]
pub struct UnionFind {
    inner: Vec<Node>,
}

コンストラクタ：new()

Union Find Tree は独立したサイズ１の木からなる森として初期化されます。

はじめは全員おひとり様グループです。ここに友達関係を導入していきます。

pub fn new(size: usize) -> Self {
    Self {
        inner: vec![Node::Size(1); size],
    }
}

ルートインデックスを得る：find_root_idx()

i番目の要素が所属するグループのルートのインデックスを得る関数です。i番目の要素がルートであればiを返します。そうでない場合は、ルートまで親をたどっていきます。この時に親のインデックスをルートのインデックスに更新しています。これにより、木が浅くなり次回以降の計算コストを削減できます。このテクニックを経路圧縮といいます。経路圧縮により部分木の高さが１になります。

経路圧縮のイメージ

find_root_idx(d)を実行すると木の構造が変化します。例えばfind_root_idx(c)の計算コストが削減されています。

pub fn find_root_idx(&mut self, i: usize) -> usize {
    assert!(i < self.inner.len());

    match self.inner[i] {
        Node::Size(_) => i,
        Node::Parent(parent_idx) => {
            let root_idx = self.find_root_idx(parent_idx); // path compression
            self.inner[i] = Node::Parent(root_idx);
            root_idx
        }
    }
}

２要素が同じグループに所属しているか：same()

find_root_idx()を利用して、ルートインデックスを比較します。

pub fn same(&mut self, i: usize, j: usize) -> bool {
    assert!(i < self.inner.len());
    assert!(j < self.inner.len());

    self.find_root_idx(i) == self.find_root_idx(j)
}

グループのサイズを得る：get_size()

find_root_idx()と同じようにして、グループのサイズを得ることができます。

pub fn get_size(&mut self, i: usize) -> usize {
    assert!(i < self.inner.len());

    match self.inner[i] {
        Node::Size(size) => size,
        Node::Parent(parent_idx) => self.find_root_idx(parent_idx),
    }
}

２つの要素をグルーピングする：unite()

ここまでで、２つのグループを結合する操作が実現できます。

  1. i番目の要素とj番目の要素のルートインデックスを比較します。同じなら操作を終了し、異なる場合は次のステップに移ります。
  2. サイズが小さいグループを大きいグループに結合します。これをunion by sizeといいます。これによって、木の深さをできるだけ浅くなるようにします。
pub fn unite(&mut self, i: usize, j: usize) -> bool {
    assert!(i < self.inner.len());
    assert!(j < self.inner.len());

    let mut ri = self.find_root_idx(i);
    let mut rj = self.find_root_idx(j);
    if ri == rj {
        return false;
    }

    // union by size
    if self.get_size(ri) < self.get_size(rj) {
        std::mem::swap(&mut ri, &mut rj);
    }

    self.inner[ri] = Node::Size(self.get_size(ri) + self.get_size(rj));
    self.inner[rj] = Node::Parent(ri);

    true
}

グループサイズの一覧を得る：size_iter()

オプションです。実装上のポイントはcollect()関数は呼び出し側に任せることです。これによって、vec.into_iter().collect().into_iter()のような無駄な処理を省くことができます。イテレーターはイテレーターのまま返しましょう。

pub fn size_iter(&self) -> impl Iterator<Item = usize> + '_ {
    self.inner.iter().filter_map(|v| match v {
        Node::Size(size) => Some(*size),
        Node::Parent(_) => None,
    })
}
脚注

  1. ほとんど定数時間。 ↩︎

Discussion