Home > libalgo > Union find tree

Union find tree

概要

Union find.

実装

#[allow(dead_code)]
struct UnionFind {
    par: Vec<i32>,
    sz: usize,
}

#[allow(dead_code)]
impl UnionFind {
    fn new(n: usize) -> UnionFind {
        UnionFind {
            par: vec![-1; n],
            sz: n,
        }
    }

    fn find_root(&mut self, a: usize) -> usize {
        if self.par[a] < 0 {
            a
        } else {
            let t = self.par[a];
            self.par[a] = self.find_root(t as usize) as i32;
            self.par[a] as usize
        }
    }

    fn merge(&mut self, a: usize, b: usize) -> bool {
        let a = self.find_root(a);
        let b = self.find_root(b);
        if a == b {
            false
        } else {
            let (a, b) = if self.par[a] > self.par[b] {
                (b, a)
            } else {
                (a, b)
            };
            self.par[a] += self.par[b];
            self.par[b] = a as i32;
            self.sz -= 1;
            true
        }
    }

    fn is_same(&mut self, a: usize, b: usize) -> bool {
        self.find_root(a) == self.find_root(b)
    }

    fn size(&mut self, a: usize) -> usize {
        let a = self.find_root(a);
        -self.par[a] as usize
    }

    fn count_groups(&self) -> usize {
        self.sz
    }
}