geop_core_math/
union_find.rs

1//! [`UnionFind`]: which of `n` things have been joined into one group —
2//! connected faces of a shell, the cells a split cuts a solid into.
3
4/// Groups of the indices `0..n`, merged by [`UnionFind::union`].
5#[derive(Clone, Debug)]
6pub struct UnionFind {
7    parent: Vec<usize>,
8}
9
10impl UnionFind {
11    /// `n` indices, each in a group of its own.
12    pub fn new(n: usize) -> Self {
13        Self {
14            parent: (0..n).collect(),
15        }
16    }
17
18    /// The representative of `i`'s group: the same for every member.
19    pub fn find(&mut self, i: usize) -> usize {
20        let mut root = i;
21        while self.parent[root] != root {
22            root = self.parent[root];
23        }
24        let mut at = i;
25        while self.parent[at] != root {
26            at = std::mem::replace(&mut self.parent[at], root);
27        }
28        root
29    }
30
31    /// Merges the groups of `a` and `b`.
32    pub fn union(&mut self, a: usize, b: usize) {
33        let (a, b) = (self.find(a), self.find(b));
34        if a != b {
35            self.parent[a.max(b)] = a.min(b);
36        }
37    }
38
39    /// Every group, each in increasing order, ordered by its smallest
40    /// member.
41    pub fn groups(&mut self) -> Vec<Vec<usize>> {
42        let mut by_root: std::collections::BTreeMap<usize, Vec<usize>> = Default::default();
43        for i in 0..self.parent.len() {
44            let root = self.find(i);
45            by_root.entry(root).or_default().push(i);
46        }
47        let mut groups: Vec<Vec<usize>> = by_root.into_values().collect();
48        groups.sort_by_key(|g| g[0]);
49        groups
50    }
51}
52
53#[cfg(test)]
54mod tests {
55    use super::UnionFind;
56
57    #[test]
58    fn groups_follow_unions() {
59        let mut sets = UnionFind::new(6);
60        sets.union(4, 1);
61        sets.union(2, 5);
62        sets.union(5, 4);
63        assert_eq!(sets.groups(), vec![vec![0], vec![1, 2, 4, 5], vec![3]]);
64        assert_eq!(sets.find(2), sets.find(1));
65    }
66}
geop_core_math/union_find.rs

geop_core_math/
union_find.rs
