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}