Skip to content

Commit a440df7

Browse files
sozelfistsozelfist
committed
ref: refactor UnionFind
1 parent 4e59645 commit a440df7

File tree

1 file changed

+138
-133
lines changed

1 file changed

+138
-133
lines changed

‎src/data_structures/union_find.rs

Lines changed: 138 additions & 133 deletions
Original file line numberDiff line numberDiff line change
@@ -1,27 +1,24 @@
1+
//! A Union-Find (Disjoint Set) data structure implementation in Rust.
2+
//!
3+
//! The Union-Find data structure keeps track of elements partitioned into
4+
//! disjoint (non-overlapping) sets.
5+
//! It provides near-constant-time operations to add new sets, to find the
6+
//! representative of a set, and to merge sets.
7+
18
use std::collections::HashMap;
29
use std::fmt::Debug;
310
use std::hash::Hash;
411

5-
/// UnionFind data structure
6-
/// It acts by holding an array of pointers to parents, together with the size of each subset
712
#[derive(Debug)]
813
pub struct UnionFind<T: Debug + Eq + Hash> {
9-
payloads: HashMap<T, usize>, // we are going to manipulate indices to parent, thus `usize`. We need a map to associate a value to its index in the parent links array
10-
parent_links: Vec<usize>, // holds the relationship between an item and its parent. The root of a set is denoted by parent_links[i] == i
11-
sizes: Vec<usize>, // holds the size
12-
count: usize,
14+
payloads: HashMap<T, usize>, // Maps values to their indices in the parent_links array.
15+
parent_links: Vec<usize>, // Holds the parent pointers; root elements are their own parents.
16+
sizes: Vec<usize>, // Holds the sizes of the sets.
17+
count: usize, // Number of disjoint sets.
1318
}
1419

1520
impl<T: Debug + Eq + Hash> UnionFind<T> {
16-
/// Creates an empty Union Find structure with capacity n
17-
///
18-
/// # Examples
19-
///
20-
/// ```
21-
/// use the_algorithms_rust::data_structures::UnionFind;
22-
/// let uf = UnionFind::<&str>::with_capacity(5);
23-
/// assert_eq!(0, uf.count())
24-
/// ```
21+
/// Creates an empty Union-Find structure with a specified capacity.
2522
pub fn with_capacity(capacity: usize) -> Self {
2623
Self {
2724
parent_links: Vec::with_capacity(capacity),
@@ -31,7 +28,7 @@ impl<T: Debug + Eq + Hash> UnionFind<T> {
3128
}
3229
}
3330

34-
/// Inserts a new item (disjoint) in the data structure
31+
/// Inserts a new item (disjoint set) into the data structure.
3532
pub fn insert(&mut self, item: T) {
3633
let key = self.payloads.len();
3734
self.parent_links.push(key);
@@ -40,107 +37,60 @@ impl<T: Debug + Eq + Hash> UnionFind<T> {
4037
self.count += 1;
4138
}
4239

43-
pub fn id(&self, value: &T) -> Option<usize> {
44-
self.payloads.get(value).copied()
45-
}
46-
47-
/// Returns the key of an item stored in the data structure or None if it doesn't exist
48-
fn find(&self, value: &T) -> Option<usize> {
49-
self.id(value).map(|id| self.find_by_key(id))
40+
/// Returns the root index of the set containing the given value, or `None` if it doesn't exist.
41+
pub fn find(&mut self, value: &T) -> Option<usize> {
42+
self.payloads
43+
.get(value)
44+
.copied()
45+
.map(|key| self.find_by_key(key))
5046
}
5147

52-
/// Creates a link between value_1 and value_2
53-
/// returns None if either value_1 or value_2 hasn't been inserted in the data structure first
54-
/// returns Some(true) if two disjoint sets have been merged
55-
/// returns Some(false) if both elements already were belonging to the same set
56-
///
57-
/// #_Examples:
58-
///
59-
/// ```
60-
/// use the_algorithms_rust::data_structures::UnionFind;
61-
/// let mut uf = UnionFind::with_capacity(2);
62-
/// uf.insert("A");
63-
/// uf.insert("B");
64-
///
65-
/// assert_eq!(None, uf.union(&"A", &"C"));
66-
///
67-
/// assert_eq!(2, uf.count());
68-
/// assert_eq!(Some(true), uf.union(&"A", &"B"));
69-
/// assert_eq!(1, uf.count());
70-
///
71-
/// assert_eq!(Some(false), uf.union(&"A", &"B"));
72-
/// ```
73-
pub fn union(&mut self, item1: &T, item2: &T) -> Option<bool> {
74-
match (self.find(item1), self.find(item2)) {
75-
(Some(k1), Some(k2)) => Some(self.union_by_key(k1, k2)),
48+
/// Unites the sets containing the two given values. Returns:
49+
/// - `None` if either value hasn't been inserted,
50+
/// - `Some(true)` if two disjoint sets have been merged,
51+
/// - `Some(false)` if both elements were already in the same set.
52+
pub fn union(&mut self, first_item: &T, sec_item: &T) -> Option<bool> {
53+
let (first_root, sec_root) = (self.find(first_item), self.find(sec_item));
54+
match (first_root, sec_root) {
55+
(Some(first_root), Some(sec_root)) => Some(self.union_by_key(first_root, sec_root)),
7656
_ => None,
7757
}
7858
}
7959

80-
/// Returns the parent of the element given its id
81-
fn find_by_key(&self, key: usize) -> usize {
82-
let mut id = key;
83-
while id != self.parent_links[id] {
84-
id = self.parent_links[id];
60+
/// Finds the root of the set containing the element with the given index.
61+
fn find_by_key(&mut self, key: usize) -> usize {
62+
if self.parent_links[key] != key {
63+
self.parent_links[key] = self.find_by_key(self.parent_links[key]);
8564
}
86-
id
65+
self.parent_links[key]
8766
}
8867

89-
/// Unions the sets containing id1 and id2
90-
fn union_by_key(&mut self, key1: usize, key2: usize) -> bool {
91-
let root1 = self.find_by_key(key1);
92-
let root2 = self.find_by_key(key2);
93-
if root1 == root2 {
94-
return false; // they belong to the same set already, no-op
68+
/// Unites the sets containing the two elements identified by their indices.
69+
fn union_by_key(&mut self, first_key: usize, sec_key: usize) -> bool {
70+
let (first_root, sec_root) = (self.find_by_key(first_key), self.find_by_key(sec_key));
71+
72+
if first_root == sec_root {
73+
return false;
9574
}
96-
// Attach the smaller set to the larger one
97-
if self.sizes[root1] < self.sizes[root2] {
98-
self.parent_links[root1] = root2;
99-
self.sizes[root2] += self.sizes[root1];
75+
76+
if self.sizes[first_root] < self.sizes[sec_root] {
77+
self.parent_links[first_root] = sec_root;
78+
self.sizes[sec_root] += self.sizes[first_root];
10079
} else {
101-
self.parent_links[root2] = root1;
102-
self.sizes[root1] += self.sizes[root2];
80+
self.parent_links[sec_root] = first_root;
81+
self.sizes[first_root] += self.sizes[sec_root];
10382
}
104-
self.count -= 1; // we had 2 disjoint sets, now merged as one
83+
84+
self.count -= 1;
10585
true
10686
}
10787

108-
/// Checks if two items belong to the same set
109-
///
110-
/// #_Examples:
111-
///
112-
/// ```
113-
/// use the_algorithms_rust::data_structures::UnionFind;
114-
/// let mut uf = UnionFind::from_iter(["A", "B"]);
115-
/// assert!(!uf.is_same_set(&"A", &"B"));
116-
///
117-
/// uf.union(&"A", &"B");
118-
/// assert!(uf.is_same_set(&"A", &"B"));
119-
///
120-
/// assert!(!uf.is_same_set(&"A", &"C"));
121-
/// ```
122-
pub fn is_same_set(&self, item1: &T, item2: &T) -> bool {
123-
matches!((self.find(item1), self.find(item2)), (Some(root1), Some(root2)) if root1 == root2)
88+
/// Checks if two items belong to the same set.
89+
pub fn is_same_set(&mut self, first_item: &T, sec_item: &T) -> bool {
90+
matches!((self.find(first_item), self.find(sec_item)), (Some(first_root), Some(sec_root)) if first_root == sec_root)
12491
}
12592

126-
/// Returns the number of disjoint sets
127-
///
128-
/// # Examples
129-
///
130-
/// ```
131-
/// use the_algorithms_rust::data_structures::UnionFind;
132-
/// let mut uf = UnionFind::with_capacity(5);
133-
/// assert_eq!(0, uf.count());
134-
///
135-
/// uf.insert("A");
136-
/// assert_eq!(1, uf.count());
137-
///
138-
/// uf.insert("B");
139-
/// assert_eq!(2, uf.count());
140-
///
141-
/// uf.union(&"A", &"B");
142-
/// assert_eq!(1, uf.count())
143-
/// ```
93+
/// Returns the number of disjoint sets.
14494
pub fn count(&self) -> usize {
14595
self.count
14696
}
@@ -158,11 +108,11 @@ impl<T: Debug + Eq + Hash> Default for UnionFind<T> {
158108
}
159109

160110
impl<T: Debug + Eq + Hash> FromIterator<T> for UnionFind<T> {
161-
/// Creates a new UnionFind data structure from an iterable of disjoint elements
111+
/// Creates a new UnionFind data structure from an iterable of disjoint elements.
162112
fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
163113
let mut uf = UnionFind::default();
164-
for i in iter {
165-
uf.insert(i);
114+
for item in iter {
115+
uf.insert(item);
166116
}
167117
uf
168118
}
@@ -175,45 +125,100 @@ mod tests {
175125
#[test]
176126
fn test_union_find() {
177127
let mut uf = UnionFind::from_iter(0..10);
178-
assert_eq!(uf.find_by_key(0), 0);
179-
assert_eq!(uf.find_by_key(1), 1);
180-
assert_eq!(uf.find_by_key(2), 2);
181-
assert_eq!(uf.find_by_key(3), 3);
182-
assert_eq!(uf.find_by_key(4), 4);
183-
assert_eq!(uf.find_by_key(5), 5);
184-
assert_eq!(uf.find_by_key(6), 6);
185-
assert_eq!(uf.find_by_key(7), 7);
186-
assert_eq!(uf.find_by_key(8), 8);
187-
assert_eq!(uf.find_by_key(9), 9);
188-
189-
assert_eq!(Some(true), uf.union(&0, &1));
190-
assert_eq!(Some(true), uf.union(&1, &2));
191-
assert_eq!(Some(true), uf.union(&2, &3));
128+
assert_eq!(uf.find(&0), Some(0));
129+
assert_eq!(uf.find(&1), Some(1));
130+
assert_eq!(uf.find(&2), Some(2));
131+
assert_eq!(uf.find(&3), Some(3));
132+
assert_eq!(uf.find(&4), Some(4));
133+
assert_eq!(uf.find(&5), Some(5));
134+
assert_eq!(uf.find(&6), Some(6));
135+
assert_eq!(uf.find(&7), Some(7));
136+
assert_eq!(uf.find(&8), Some(8));
137+
assert_eq!(uf.find(&9), Some(9));
138+
139+
assert!(!uf.is_same_set(&0, &1));
140+
assert!(!uf.is_same_set(&2, &9));
141+
assert_eq!(uf.count(), 10);
142+
143+
assert_eq!(uf.union(&0, &1), Some(true));
144+
assert_eq!(uf.union(&1, &2), Some(true));
145+
assert_eq!(uf.union(&2, &3), Some(true));
146+
assert_eq!(uf.union(&0, &2), Some(false));
147+
assert_eq!(uf.union(&4, &5), Some(true));
148+
assert_eq!(uf.union(&5, &6), Some(true));
149+
assert_eq!(uf.union(&6, &7), Some(true));
150+
assert_eq!(uf.union(&7, &8), Some(true));
151+
assert_eq!(uf.union(&8, &9), Some(true));
152+
assert_eq!(uf.union(&7, &9), Some(false));
153+
154+
assert_ne!(uf.find(&0), uf.find(&9));
155+
assert_eq!(uf.find(&0), uf.find(&3));
156+
assert_eq!(uf.find(&4), uf.find(&9));
157+
assert!(uf.is_same_set(&0, &3));
158+
assert!(uf.is_same_set(&4, &9));
159+
assert!(!uf.is_same_set(&0, &9));
160+
assert_eq!(uf.count(), 2);
161+
192162
assert_eq!(Some(true), uf.union(&3, &4));
193-
assert_eq!(Some(true), uf.union(&4, &5));
194-
assert_eq!(Some(true), uf.union(&5, &6));
195-
assert_eq!(Some(true), uf.union(&6, &7));
196-
assert_eq!(Some(true), uf.union(&7, &8));
197-
assert_eq!(Some(true), uf.union(&8, &9));
198-
assert_eq!(Some(false), uf.union(&9, &0));
199-
200-
assert_eq!(1, uf.count());
163+
assert_eq!(uf.find(&0), uf.find(&9));
164+
assert_eq!(uf.count(), 1);
165+
assert!(uf.is_same_set(&0, &9));
166+
167+
assert_eq!(None, uf.union(&0, &11));
201168
}
202169

203170
#[test]
204171
fn test_spanning_tree() {
205-
// Let's imagine the following topology:
206-
// A <-> B
207-
// B <-> C
208-
// A <-> D
209-
// E
210-
// F <-> G
211-
// We have 3 disjoint sets: {A, B, C, D}, {E}, {F, G}
212172
let mut uf = UnionFind::from_iter(["A", "B", "C", "D", "E", "F", "G"]);
213173
uf.union(&"A", &"B");
214174
uf.union(&"B", &"C");
215175
uf.union(&"A", &"D");
216176
uf.union(&"F", &"G");
217-
assert_eq!(3, uf.count());
177+
178+
assert_eq!(None, uf.union(&"A", &"W"));
179+
180+
assert_eq!(uf.find(&"A"), uf.find(&"B"));
181+
assert_eq!(uf.find(&"A"), uf.find(&"C"));
182+
assert_eq!(uf.find(&"B"), uf.find(&"D"));
183+
assert_ne!(uf.find(&"A"), uf.find(&"E"));
184+
assert_ne!(uf.find(&"A"), uf.find(&"F"));
185+
assert_eq!(uf.find(&"G"), uf.find(&"F"));
186+
assert_ne!(uf.find(&"G"), uf.find(&"E"));
187+
188+
assert!(uf.is_same_set(&"A", &"B"));
189+
assert!(uf.is_same_set(&"A", &"C"));
190+
assert!(uf.is_same_set(&"B", &"D"));
191+
assert!(!uf.is_same_set(&"B", &"F"));
192+
assert!(!uf.is_same_set(&"E", &"A"));
193+
assert!(!uf.is_same_set(&"E", &"G"));
194+
assert_eq!(uf.count(), 3);
195+
}
196+
197+
#[test]
198+
fn test_with_capacity() {
199+
let mut uf: UnionFind<i32> = UnionFind::with_capacity(5);
200+
uf.insert(0);
201+
uf.insert(1);
202+
uf.insert(2);
203+
uf.insert(3);
204+
uf.insert(4);
205+
206+
assert_eq!(uf.count(), 5);
207+
208+
assert_eq!(uf.union(&0, &1), Some(true));
209+
assert!(uf.is_same_set(&0, &1));
210+
assert_eq!(uf.count(), 4);
211+
212+
assert_eq!(uf.union(&2, &3), Some(true));
213+
assert!(uf.is_same_set(&2, &3));
214+
assert_eq!(uf.count(), 3);
215+
216+
assert_eq!(uf.union(&0, &2), Some(true));
217+
assert!(uf.is_same_set(&0, &1));
218+
assert!(uf.is_same_set(&2, &3));
219+
assert!(uf.is_same_set(&0, &3));
220+
assert_eq!(uf.count(), 2);
221+
222+
assert_eq!(None, uf.union(&0, &10));
218223
}
219224
}

0 commit comments

Comments
 (0)