dsalgo 0.3.10

A package for Datastructures and Algorithms.
Documentation
pub struct PotentialUnionFind {
    a: Vec<isize>,
    rh: Vec<i64>,
}

impl PotentialUnionFind {
    pub fn size(&self) -> usize {
        self.a.len()
    }

    pub fn new(size: usize) -> Self {
        Self { a: vec![-1; size], rh: vec![0; size] }
    }

    pub fn root(
        &mut self,
        u: usize,
    ) -> usize {
        if self.a[u] < 0 {
            return u;
        }

        let p = self.a[u] as usize;

        self.a[u] = self.root(p) as isize;

        self.rh[u] += self.rh[p];

        self.a[u] as usize
    }

    fn h(
        &mut self,
        u: usize,
    ) -> i64 {
        self.root(u);

        self.rh[u]
    }

    pub fn size_of(
        &mut self,
        u: usize,
    ) -> usize {
        let u = self.root(u);

        -self.a[u] as usize
    }

    pub fn same(
        &mut self,
        u: usize,
        v: usize,
    ) -> bool {
        self.root(u) == self.root(v)
    }

    pub fn diff(
        &mut self,
        u: usize,
        v: usize,
    ) -> Result<i64, &'static str> {
        if !self.same(u, v) {
            Err("belongs to different components")
        } else {
            Ok(self.h(v) - self.h(u))
        }
    }

    pub fn unite(
        &mut self,
        mut u: usize,
        mut v: usize,
        diff: i64,
    ) {
        let mut d = self.h(u) + diff - self.h(v);

        u = self.root(u);

        v = self.root(v);

        if u == v {
            debug_assert!(d == 0);

            return;
        }

        if self.a[u] > self.a[v] {
            std::mem::swap(&mut u, &mut v);

            d = -d;
        }

        self.a[u] += self.a[v];

        self.a[v] = u as isize;

        self.rh[v] = d;
    }
}

#[cfg(test)]

mod tests {

    use super::*;

    #[test]

    fn test_potential_uf() {
        let mut uf = PotentialUnionFind::new(6);

        assert_eq!(uf.size_of(0), 1);

        assert!(uf.diff(0, 5).is_err());

        uf.unite(0, 1, 1);

        uf.unite(5, 4, 2);

        uf.unite(1, 2, 3);

        uf.unite(4, 3, 4);

        uf.unite(2, 3, 5);

        assert_eq!(uf.size_of(4), 6);

        assert_eq!(uf.diff(0, 5), Ok(3));
    }
}