graph_spanning/disjoint_set.rs
1/// Union-Find (Disjoint Set Union) with **union by rank** and **path
2/// compression**.
3///
4/// This is the complete version of the skeleton from `graph-collections`.
5/// Path compression flattens the tree during every [`Self::find`] call, giving
6/// amortised O(α(n)) per operation where α is the inverse Ackermann function
7/// — effectively constant for all practical inputs.
8///
9/// # Examples
10///
11/// ```
12/// use graph_spanning::DisjointSet;
13///
14/// let mut ds = DisjointSet::new(5);
15/// assert_eq!(ds.count(), 5);
16///
17/// ds.union(0, 1);
18/// ds.union(2, 3);
19/// assert!(ds.connected(0, 1));
20/// assert!(ds.connected(2, 3));
21/// assert!(!ds.connected(0, 2));
22/// assert_eq!(ds.count(), 3);
23///
24/// ds.union(1, 3); // merges the two groups
25/// assert!(ds.connected(0, 3));
26/// assert_eq!(ds.count(), 2);
27/// ```
28#[derive(Debug, Clone)]
29pub struct DisjointSet {
30 parent: Vec<usize>,
31 rank: Vec<usize>,
32 count: usize,
33}
34
35impl DisjointSet {
36 /// Creates a new Union-Find with `n` elements, each in its own set.
37 ///
38 /// Elements are indexed `0..n`.
39 ///
40 /// # Examples
41 ///
42 /// ```
43 /// use graph_spanning::DisjointSet;
44 ///
45 /// let ds = DisjointSet::new(4);
46 /// assert_eq!(ds.count(), 4);
47 /// ```
48 pub fn new(n: usize) -> Self {
49 Self {
50 parent: (0..n).collect(),
51 rank: vec![0; n],
52 count: n,
53 }
54 }
55
56 /// Finds the representative (root) of the set containing `x`.
57 ///
58 /// Applies **path compression**: every node on the path to the root is
59 /// re-pointed directly to the root, flattening the tree for future calls.
60 ///
61 /// # Panics
62 ///
63 /// Panics if `x >= n` (out of bounds).
64 ///
65 /// # Examples
66 ///
67 /// ```
68 /// use graph_spanning::DisjointSet;
69 ///
70 /// let mut ds = DisjointSet::new(3);
71 /// ds.union(0, 1);
72 /// // After union, both 0 and 1 share the same representative.
73 /// assert_eq!(ds.find(0), ds.find(1));
74 /// ```
75 pub fn find(&mut self, x: usize) -> usize {
76 if self.parent[x] != x {
77 // Recursive path compression: point x directly to the root.
78 self.parent[x] = self.find(self.parent[x]);
79 }
80 self.parent[x]
81 }
82
83 /// Merges the sets containing `x` and `y`.
84 ///
85 /// Uses **union by rank**: the root of the smaller-rank tree is attached
86 /// to the root of the larger-rank tree, keeping the tree shallow.
87 ///
88 /// Returns `true` if `x` and `y` were in different sets (a merge
89 /// occurred), or `false` if they were already in the same set.
90 ///
91 /// # Panics
92 ///
93 /// Panics if `x` or `y` is out of bounds.
94 ///
95 /// # Examples
96 ///
97 /// ```
98 /// use graph_spanning::DisjointSet;
99 ///
100 /// let mut ds = DisjointSet::new(3);
101 /// assert!(ds.union(0, 1)); // merged
102 /// assert!(!ds.union(0, 1)); // already connected
103 /// ```
104 pub fn union(&mut self, x: usize, y: usize) -> bool {
105 let rx = self.find(x);
106 let ry = self.find(y);
107
108 if rx == ry {
109 return false; // already in the same set
110 }
111
112 // Attach smaller-rank tree under larger-rank root.
113 match self.rank[rx].cmp(&self.rank[ry]) {
114 std::cmp::Ordering::Less => self.parent[rx] = ry,
115 std::cmp::Ordering::Greater => self.parent[ry] = rx,
116 std::cmp::Ordering::Equal => {
117 self.parent[ry] = rx;
118 self.rank[rx] += 1;
119 }
120 }
121
122 self.count -= 1;
123 true
124 }
125
126 /// Returns `true` if `x` and `y` are in the same set.
127 ///
128 /// # Panics
129 ///
130 /// Panics if `x` or `y` is out of bounds.
131 ///
132 /// # Examples
133 ///
134 /// ```
135 /// use graph_spanning::DisjointSet;
136 ///
137 /// let mut ds = DisjointSet::new(4);
138 /// ds.union(0, 1);
139 /// ds.union(1, 2);
140 /// assert!(ds.connected(0, 2)); // transitive
141 /// assert!(!ds.connected(0, 3));
142 /// ```
143 pub fn connected(&mut self, x: usize, y: usize) -> bool {
144 self.find(x) == self.find(y)
145 }
146
147 /// Returns the current number of disjoint sets.
148 ///
149 /// Starts at `n` and decreases by 1 for every successful [`Self::union`] call.
150 ///
151 /// # Examples
152 ///
153 /// ```
154 /// use graph_spanning::DisjointSet;
155 ///
156 /// let mut ds = DisjointSet::new(4);
157 /// assert_eq!(ds.count(), 4);
158 /// ds.union(0, 1);
159 /// assert_eq!(ds.count(), 3);
160 /// ```
161 #[inline]
162 pub fn count(&self) -> usize {
163 self.count
164 }
165}