haar_lib/ds/
potential_unionfind.rs

1//! ポテンシャル付きUnionfind
2//!
3//! # Problems
4//! - <https://judge.yosupo.jp/problem/unionfind_with_potential>
5//! - <https://judge.yosupo.jp/problem/unionfind_with_potential_non_commutative_group>
6
7use crate::algebra::traits::*;
8use std::cell::{Cell, RefCell};
9
10/// ポテンシャル付きUnionfind
11pub struct PotentialUnionFind<G: Group> {
12    group: G,
13    n: usize,
14    count: usize,
15    parent: Vec<Cell<usize>>,
16    depth: Vec<usize>,
17    size: Vec<usize>,
18    potential: RefCell<Vec<G::Element>>,
19    inverse: Option<RefCell<Vec<G::Element>>>,
20}
21
22impl<G> PotentialUnionFind<G>
23where
24    G: AbelianGroup,
25    G::Element: Clone,
26{
27    /// 大きさ`n`の[`PotentialUnionFind`]を生成する。(ポテンシャルが可換群のとき)
28    pub fn new_commutative(group: G, n: usize) -> Self {
29        Self {
30            n,
31            count: n,
32            parent: (0..n).map(Cell::new).collect(),
33            depth: vec![1; n],
34            size: vec![1; n],
35            potential: RefCell::new(vec![group.id(); n]),
36            inverse: None,
37            group,
38        }
39    }
40}
41
42impl<G> PotentialUnionFind<G>
43where
44    G: Group,
45    G::Element: Clone,
46{
47    /// 大きさ`n`の[`PotentialUnionFind`]を生成する。(ポテンシャルが可換とは限らない群のとき)
48    pub fn new_non_commutative(group: G, n: usize) -> Self {
49        Self {
50            n,
51            count: n,
52            parent: (0..n).map(Cell::new).collect(),
53            depth: vec![1; n],
54            size: vec![1; n],
55            potential: RefCell::new(vec![group.id(); n]),
56            inverse: Some(RefCell::new(vec![group.id(); n])),
57            group,
58        }
59    }
60
61    /// `i`の属する素集合の根を返す。
62    pub fn root_of(&self, i: usize) -> usize {
63        if self.parent[i].get() == i {
64            return i;
65        }
66        let p = self.parent[i].get();
67        let p = self.root_of(p);
68
69        let mut potential = self.potential.borrow_mut();
70        let t = potential[self.parent[i].get()].clone();
71        self.group.op_assign_l(&mut potential[i], t);
72
73        if let Some(inv) = self.inverse.as_ref() {
74            let mut inv = inv.borrow_mut();
75            let t = inv[self.parent[i].get()].clone();
76            self.group.op_assign_r(&mut inv[i], t);
77        }
78
79        self.parent[i].set(p);
80        self.parent[i].get()
81    }
82
83    /// `i`のポテンシャル($P(i)$)を返す。
84    pub fn potential_of(&self, i: usize) -> G::Element {
85        self.potential.borrow()[i].clone()
86    }
87
88    /// `i`と`j`が同じ素集合に属するならば`true`を返す。
89    pub fn is_same(&self, i: usize, j: usize) -> bool {
90        self.root_of(i) == self.root_of(j)
91    }
92
93    /// `i`と`j`が同一の素集合に属するとき、ポテンシャルの差($P(i) - P(j)$)を返す。
94    pub fn diff(&self, i: usize, j: usize) -> Option<G::Element> {
95        self.is_same(i, j).then(|| {
96            let pi = self.potential_of(i);
97            if let Some(inv) = self.inverse.as_ref() {
98                self.group.op(inv.borrow()[j].clone(), pi)
99            } else {
100                self.group.op(self.group.inv(self.potential_of(j)), pi)
101            }
102        })
103    }
104
105    /// `i`の属する素集合と`j`の属する素集合を統合する。
106    /// 統合後は、$P(i) = P(j) + p$を満たす。
107    pub fn merge(&mut self, i: usize, j: usize, p: G::Element) -> usize {
108        let ri = self.root_of(i);
109        let rj = self.root_of(j);
110
111        if ri == rj {
112            return ri;
113        }
114
115        self.count -= 1;
116
117        let mut potential = self.potential.borrow_mut();
118        let (pi, pj) = (potential[i].clone(), potential[j].clone());
119
120        let g = &self.group;
121
122        if self.depth[ri] < self.depth[rj] {
123            self.parent[ri].set(rj);
124            self.size[rj] += self.size[ri];
125
126            if let Some(inv) = self.inverse.as_ref() {
127                let mut inv = inv.borrow_mut();
128                potential[ri] = g.fold_m([pj, p.clone(), inv[i].clone()]);
129                inv[ri] = g.fold_m([pi, g.inv(p), inv[j].clone()]);
130            } else {
131                potential[ri] = g.fold_m([pj, p, g.inv(pi)]);
132            }
133
134            j
135        } else {
136            self.parent[rj].set(ri);
137            self.size[ri] += self.size[rj];
138
139            if let Some(inv) = self.inverse.as_ref() {
140                let mut inv = inv.borrow_mut();
141                potential[rj] = g.fold_m([pi, g.inv(p.clone()), inv[j].clone()]);
142                inv[rj] = g.fold_m([pj, p, inv[i].clone()]);
143            } else {
144                potential[rj] = g.fold_m([pi, g.inv(p), g.inv(pj)]);
145            }
146
147            if self.depth[ri] == self.depth[rj] {
148                self.depth[i] += 1;
149            }
150
151            i
152        }
153    }
154
155    /// `i`の属する素集合の大きさを返す。
156    pub fn size_of(&self, i: usize) -> usize {
157        self.size[self.root_of(i)]
158    }
159
160    /// 素集合の個数を返す。
161    pub fn count_groups(&self) -> usize {
162        self.count
163    }
164
165    /// 素集合をすべて列挙する。
166    pub fn get_groups(&self) -> Vec<Vec<usize>> {
167        let mut ret = vec![vec![]; self.n];
168
169        for i in 0..self.n {
170            ret[self.root_of(i)].push(i);
171        }
172
173        ret.into_iter().filter(|x| !x.is_empty()).collect()
174    }
175}
haar_lib/ds/potential_unionfind.rs

haar_lib/ds/
potential_unionfind.rs
