data_beans/alg/
union_find.rs

1//! Disjoint-set union (DSU) with union-by-rank and path halving.
2//!
3//! Kept in-tree (not delegated to `petgraph`/`disjoint_sets`) because callers
4//! need the per-root **size** (member count) for weighted centroid merging — a
5//! field none of the ecosystem crates expose.
6
7/// Union-Find (disjoint set) with path halving and union by rank.
8pub struct UnionFind {
9    parent: Vec<usize>,
10    rank: Vec<usize>,
11    size: Vec<usize>,
12}
13
14impl UnionFind {
15    pub fn new(n: usize) -> Self {
16        Self {
17            parent: (0..n).collect(),
18            rank: vec![0; n],
19            size: vec![1; n],
20        }
21    }
22
23    #[inline]
24    pub fn find(&mut self, mut x: usize) -> usize {
25        while self.parent[x] != x {
26            self.parent[x] = self.parent[self.parent[x]];
27            x = self.parent[x];
28        }
29        x
30    }
31
32    /// Union the sets containing `a` and `b`; returns the surviving root.
33    #[inline]
34    pub fn union(&mut self, a: usize, b: usize) -> usize {
35        let ra = self.find(a);
36        let rb = self.find(b);
37        if ra == rb {
38            return ra;
39        }
40        let (big, small) = if self.rank[ra] >= self.rank[rb] {
41            (ra, rb)
42        } else {
43            (rb, ra)
44        };
45        self.parent[small] = big;
46        self.size[big] += self.size[small];
47        if self.rank[big] == self.rank[small] {
48            self.rank[big] += 1;
49        }
50        big
51    }
52
53    /// Member count of the set whose root is `x` (caller passes a root, typically
54    /// `uf.find(...)` — accepting any node index would require an extra find).
55    #[inline]
56    pub fn size(&self, x: usize) -> usize {
57        self.size[x]
58    }
59
60    /// Direct parent pointer — valid as a final representative only after
61    /// `flatten` has been called. Used in multilevel extraction hot loops to
62    /// avoid per-call `find` after a batch of unions.
63    #[inline]
64    pub fn parent(&self, x: usize) -> usize {
65        self.parent[x]
66    }
67
68    /// Path-compress every node so `parent[i]` is its final representative.
69    /// Call after a batch of unions if subsequent lookups use [`Self::parent`] directly.
70    pub fn flatten(&mut self) {
71        for i in 0..self.parent.len() {
72            self.parent[i] = self.find(i);
73        }
74    }
75}
data_beans/alg/union_find.rs

data_beans/alg/
union_find.rs
