nekolib/ds/
potentialized_union_find.rs

1//! ポテンシャルつき union-find。
2
3use super::super::traits::binop;
4use super::super::traits::potential_function;
5
6use std::cell::RefCell;
7
8use binop::{CommutativeGroup, Magma};
9use potential_function::PotentialFunction;
10
11#[derive(Clone, Copy)]
12enum Item {
13    Parent(usize),
14    Size(usize),
15}
16
17/// ポテンシャルつき union-find。
18///
19/// # Idea
20/// 通常の union-find に加え、配列 `pot` を管理する。
21/// 親ノード `par` と子ノード `child` に対して、`pot[child]` には
22/// `phi(child) - phi(par)` を持つようにする。
23///
24/// 代表元を探してパスを縮約する際、ポテンシャル差の更新を適切に行う。
25#[derive(Clone)]
26pub struct PotentializedUnionFind<G: CommutativeGroup>
27where
28    <G as Magma>::Set: Clone,
29{
30    n: usize,
31    buf: RefCell<Vec<Item>>,
32    pot: RefCell<Vec<<G as Magma>::Set>>,
33    cgroup: G,
34}
35
36impl<G: CommutativeGroup> PotentialFunction for PotentializedUnionFind<G>
37where
38    <G as Magma>::Set: Clone,
39{
40    type Item = G;
41    fn new(n: usize, cgroup: G) -> Self {
42        Self {
43            n,
44            buf: RefCell::new(vec![Item::Size(1); n]),
45            pot: RefCell::new(vec![cgroup.id(); n]),
46            cgroup,
47        }
48    }
49
50    fn len(&self) -> usize { self.n }
51
52    fn relate(
53        &mut self,
54        u: usize,
55        v: usize,
56        w: G::Set,
57    ) -> Result<bool, G::Set> {
58        let ru = self.repr(u);
59        let rv = self.repr(v);
60        let mut buf = self.buf.borrow_mut();
61        let mut pot = self.pot.borrow_mut();
62        // w += p[v] - p[u];
63        let diff =
64            self.cgroup.op(pot[v].clone(), self.cgroup.recip(pot[u].clone()));
65        if ru == rv {
66            let w_old = self.cgroup.recip(diff);
67            return if w == w_old { Ok(false) } else { Err(w_old) };
68        }
69        let w = self.cgroup.op(w, diff);
70
71        let (su, sv) = match (buf[ru], buf[rv]) {
72            (Item::Size(su), Item::Size(sv)) => (su, sv),
73            _ => unreachable!(),
74        };
75
76        let (child, par, d) =
77            if su < sv { (ru, rv, w) } else { (rv, ru, self.cgroup.recip(w)) };
78        buf[par] = Item::Size(su + sv);
79        buf[child] = Item::Parent(par);
80        pot[child] = d;
81        Ok(true)
82    }
83
84    fn diff(&self, u: usize, v: usize) -> Option<G::Set> {
85        if self.repr(u) != self.repr(v) {
86            return None;
87        }
88        let pot = self.pot.borrow();
89        Some(self.cgroup.op(pot[u].clone(), self.cgroup.recip(pot[v].clone())))
90    }
91
92    fn repr_diff(&self, u: usize) -> (usize, G::Set) {
93        let ru = self.repr(u);
94        (ru, self.pot.borrow()[u].clone())
95    }
96}
97
98impl<G: CommutativeGroup> PotentializedUnionFind<G>
99where
100    <G as Magma>::Set: Clone,
101{
102    pub fn with_len(n: usize) -> Self
103    where
104        G: Default,
105    {
106        Self::new(n, G::default())
107    }
108
109    fn repr(&self, mut u: usize) -> usize {
110        let mut res = u;
111        let mut buf = self.buf.borrow_mut();
112        let mut pot = self.pot.borrow_mut();
113        let mut p = self.cgroup.id();
114        while let Item::Parent(v) = buf[res] {
115            p = self.cgroup.op(p.clone(), pot[res].clone());
116            res = v;
117        }
118        let mut bu = buf[u];
119        while let Item::Parent(pu) = bu {
120            buf[u] = Item::Parent(res);
121            let tmp = p.clone();
122            p = self.cgroup.op(p.clone(), self.cgroup.recip(pot[u].clone()));
123            pot[u] = tmp;
124            u = pu;
125            bu = buf[u];
126        }
127        res
128    }
129}
130
131#[test]
132fn sanity_check() {
133    use binop::{
134        new_monoid, Associative, Commutative, Identity, Magma, PartialRecip,
135        Recip,
136    };
137
138    new_monoid! { OpXor = (u32, |x, y| x ^ y, 0, |x| x, +commutative) };
139
140    let mut uf = PotentializedUnionFind::<OpXor>::with_len(4);
141    assert_eq!(uf.relate(0, 1, 1), Ok(true));
142    assert_eq!(uf.relate(0, 2, 1), Ok(true));
143    assert_eq!(uf.relate(1, 2, 1), Err(0));
144    assert_eq!(uf.relate(1, 2, 0), Ok(false));
145
146    assert_ne!(uf.repr_diff(0).1, uf.repr_diff(1).1);
147    assert_ne!(uf.repr_diff(0).1, uf.repr_diff(2).1);
148    assert_eq!(uf.repr_diff(1).1, uf.repr_diff(2).1);
149    assert_eq!(uf.repr_diff(3), (3, 0));
150}
nekolib/ds/potentialized_union_find.rs

nekolib/ds/
potentialized_union_find.rs
