data-structure/union-find-with-potential/src/lib.rs

• View this file on GitHub
• Last update: 2025-05-29 20:17:30+09:00

Verified with

• verification/aizu-online-judge/dsl_1_b/src/main.rs
• verification/library-checker/unionfind_with_potential/src/main.rs
• verification/library-checker/unionfind_with_potential_non_commutative_group/src/main.rs

Code

#[derive(Debug, Clone)]
pub struct UnionFindWithPotential<T> {
    n: usize,
    par: Vec<usize>,
    siz: Vec<usize>,
    diff_potential: Vec<T>,
    // (Abelian) Group: operation (associative) + identity element + inverse element
    op: fn(T, T) -> T,
    e: T,
    inv: fn(T) -> T,
}

impl<T: Copy + PartialEq + Eq + Default> UnionFindWithPotential<T>
where
    T: std::ops::Neg<Output = T>,
    T: std::ops::Add<T, Output = T>,
    T: std::ops::AddAssign,
{
    pub fn new_default(n: usize) -> Self {
        fn op<T>(a: T, b: T) -> T
        where
            T: std::ops::Add<T, Output = T>,
            T: std::ops::AddAssign,
        {
            a + b
        }

        fn neg<T>(a: T) -> T
        where
            T: std::ops::Neg<Output = T>,
        {
            a.neg()
        }

        Self::new(n, op, T::default(), neg)
    }
}

impl<T: Copy + PartialEq + Eq> UnionFindWithPotential<T> {
    pub fn new(n: usize, op: fn(T, T) -> T, e: T, inv: fn(T) -> T) -> Self {
        Self {
            n,
            par: (0..n).collect::<Vec<usize>>(),
            siz: vec![1; n],
            diff_potential: vec![e; n],
            op,
            e,
            inv,
        }
    }

    pub fn merge(&mut self, x: usize, y: usize, mut w: T) -> Option<usize> {
        assert!(x < self.n);
        assert!(y < self.n);
        w = (self.op)(self.potential(x), (self.inv)((self.op)(self.potential(y), w)));
        let mut x = self.leader(x);
        let mut y = self.leader(y);
        if x == y {
            return if w == self.e { Some(x) } else { None };
        }
        if self.siz[x] < self.siz[y] {
            std::mem::swap(&mut x, &mut y);
            w = (self.inv)(w);
        }
        self.siz[x] += self.siz[y];
        self.par[y] = x;
        self.diff_potential[y] = w;
        Some(x)
    }

    pub fn same(&mut self, x: usize, y: usize) -> bool {
        assert!(x < self.n);
        assert!(y < self.n);
        self.leader(x) == self.leader(y)
    }

    pub fn leader(&mut self, x: usize) -> usize {
        assert!(x < self.n);
        if self.par[x] == x {
            return x;
        }
        let leader = self.leader(self.par[x]);
        self.diff_potential[x] = (self.op)(self.diff_potential[self.par[x]], self.diff_potential[x]);
        self.par[x] = leader;
        self.par[x]
    }

    pub fn size(&mut self, x: usize) -> usize {
        assert!(x < self.n);
        let x = self.leader(x);
        self.siz[x]
    }

    pub fn diff(&mut self, x: usize, y: usize) -> T {
        assert!(x < self.n);
        assert!(y < self.n);
        assert!(self.same(x, y));
        (self.op)((self.inv)(self.potential(y)), self.potential(x))
    }

    fn potential(&mut self, x: usize) -> T {
        self.leader(x);
        self.diff_potential[x]
    }

    pub fn groups(&mut self) -> Vec<Vec<usize>> {
        let mut res = vec![vec![]; self.n];
        for i in 0..self.n {
            res[self.leader(i)].push(i);
        }
        res.into_iter().filter(|f| !f.is_empty()).collect::<Vec<_>>()
    }
}

Traceback (most recent call last):
  File "/opt/hostedtoolcache/Python/3.11.12/x64/lib/python3.11/site-packages/onlinejudge_verify/documentation/build.py", line 71, in _render_source_code_stat
    bundled_code = language.bundle(stat.path, basedir=basedir, options={'include_paths': [basedir]}).decode()
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/hostedtoolcache/Python/3.11.12/x64/lib/python3.11/site-packages/onlinejudge_verify/languages/rust.py", line 288, in bundle
    raise NotImplementedError
NotImplementedError

Back to top page