ホーム > libalgo

Union find

概要

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(&mut self, a: usize) -> usize {
        if self.par[a] < 0 {
            a
        } else {
            let t = self.par[a];
            self.par[a] = self.find(t as usize) as i32;
            self.par[a] as usize
        }
    }

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

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

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

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