RustでUnion-Find Treeを作ってみた - プログラムのお勉強メモ

Union-Find Tree

競プロer御用達のデータ構造の1つ
- データとデータをグループ化したりする
- 同グループ内に何個のデータがあるか数え上げたりする
- グループがいくつあるか数え上げたりもできる

実装にあたり心がけたこと

今日習った trait, struct, match を使う
実装にあたっては書籍*1, *2を大いに活用させて頂いた

trait UnionFind {
    fn root(&mut self, x: isize) -> isize;
    fn issame(&mut self, x: isize, y: isize) -> bool;
    fn unite(&mut self, x: isize, y: isize);
}

#[derive(Debug)]
struct UnionFindTree {
    par: Vec<isize>,
    size: Vec<isize>,
}

impl UnionFind for UnionFindTree {
    fn root(&mut self, x: isize) -> isize {

        // 根を探す
        // 初期値(-1)なら自分が根, そうでない場合親の根に繋ぎ変えてその値を返す
        match self.par[x as usize] {
            -1 => x,
            _ => {
                self.par[x as usize] = self.root(self.par[x as usize]);
                self.par[x as usize]
            }
        }
    }

    fn issame(&mut self, x: isize, y: isize) -> bool {
        // x と y の根が同じなら true
        self.root(x) == self.root(y)
    }

    fn unite(&mut self, x: isize, y: isize) {
        let mut root_x: isize = self.root(x);
        let mut root_y: isize = self.root(y);

        // 既に同じグループなら早期リターン
        if root_x == root_y {
            return;
        }

        // 左辺の値のほうが小さかったら x と y をswap
        if self.size[root_x as usize] < self.size[root_y as usize] {
            root_x ^= root_y;
            root_y ^= root_x;
            root_x ^= root_y;
        }

        // y を x の子供として追加する
        self.par[root_y as usize] = root_x;
        self.size[root_x as usize] += self.size[root_y as usize];
    }
}

fn main() {
    let n = 7;
    let mut a = UnionFindTree {
        par: vec![-1; n],
        size: vec![1; n],
    };

    a.unite(1, 2);
    a.unite(3, 2);
    a.unite(3, 4);

    println!("{:?}", a);
    println!("{}", a.issame(1, 4));
    println!("{:?}", a);
}