Skip to content

Commit 50e67b1

Browse files
aesteveaesteve
authored
Add a payload to Union Find (TheAlgorithms#477)
1 parent da3e1f5 commit 50e67b1

File tree

1 file changed

+185
-56
lines changed

1 file changed

+185
-56
lines changed

‎src/data_structures/union_find.rs

Lines changed: 185 additions & 56 deletions
Original file line numberDiff line numberDiff line change
@@ -1,90 +1,219 @@
1+
use std::collections::HashMap;
2+
use std::fmt::Debug;
3+
use std::hash::Hash;
4+
15
/// UnionFind data structure
2-
pub struct UnionFind {
3-
id: Vec<usize>,
4-
size: Vec<usize>,
6+
/// It acts by holding an array of pointers to parents, together with the size of each subset
7+
#[derive(Debug)]
8+
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
512
count: usize,
613
}
714

8-
impl UnionFind {
9-
/// Creates a new UnionFind data structure with n elements
10-
pub fn new(n: usize) -> Self {
11-
let mut id = vec![0; n];
12-
let mut size = vec![0; n];
13-
for i in 0..n {
14-
id[i] = i;
15-
size[i] = 1;
15+
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+
/// ```
25+
pub fn with_capacity(capacity: usize) -> Self {
26+
Self {
27+
parent_links: Vec::with_capacity(capacity),
28+
sizes: Vec::with_capacity(capacity),
29+
payloads: HashMap::with_capacity(capacity),
30+
count: 0,
1631
}
17-
Self { id, size, count: n }
1832
}
1933

20-
/// Returns the parent of the element
21-
pub fn find(&mut self, x: usize) -> usize {
22-
let mut x = x;
23-
while x != self.id[x] {
24-
x = self.id[x];
25-
// self.id[x] = self.id[self.id[x]]; // path compression
34+
/// Inserts a new item (disjoint) in the data structure
35+
pub fn insert(&mut self, item: T) {
36+
let key = self.payloads.len();
37+
self.parent_links.push(key);
38+
self.sizes.push(1);
39+
self.payloads.insert(item, key);
40+
self.count += 1;
41+
}
42+
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))
50+
}
51+
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)),
76+
_ => None,
2677
}
27-
x
2878
}
2979

30-
/// Unions the sets containing x and y
31-
pub fn union(&mut self, x: usize, y: usize) -> bool {
32-
let x = self.find(x);
33-
let y = self.find(y);
34-
if x == y {
35-
return false;
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];
3685
}
37-
if self.size[x] < self.size[y] {
38-
self.id[x] = y;
39-
self.size[y] += self.size[x];
86+
id
87+
}
88+
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
95+
}
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];
40100
} else {
41-
self.id[y] = x;
42-
self.size[x] += self.size[y];
101+
self.parent_links[root2] = root1;
102+
self.sizes[root1] += self.sizes[root2];
43103
}
44-
self.count -= 1;
104+
self.count -= 1; // we had 2 disjoint sets, now merged as one
45105
true
46106
}
47107

48-
/// Checks if x and y are in the same set
49-
pub fn is_same_set(&mut self, x: usize, y: usize) -> bool {
50-
self.find(x) == self.find(y)
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)
51124
}
52125

53126
/// 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+
/// ```
54144
pub fn count(&self) -> usize {
55145
self.count
56146
}
57147
}
58148

149+
impl<T: Debug + Eq + Hash> Default for UnionFind<T> {
150+
fn default() -> Self {
151+
Self {
152+
parent_links: Vec::default(),
153+
sizes: Vec::default(),
154+
payloads: HashMap::default(),
155+
count: 0,
156+
}
157+
}
158+
}
159+
160+
impl<T: Debug + Eq + Hash> FromIterator<T> for UnionFind<T> {
161+
/// Creates a new UnionFind data structure from an iterable of disjoint elements
162+
fn from_iter<I: IntoIterator<Item = T>>(iter: I) -> Self {
163+
let mut uf = UnionFind::default();
164+
for i in iter {
165+
uf.insert(i);
166+
}
167+
uf
168+
}
169+
}
170+
59171
#[cfg(test)]
60172
mod tests {
61173
use super::*;
62174

63175
#[test]
64176
fn test_union_find() {
65-
let mut uf = UnionFind::new(10);
66-
assert_eq!(uf.find(0), 0);
67-
assert_eq!(uf.find(1), 1);
68-
assert_eq!(uf.find(2), 2);
69-
assert_eq!(uf.find(3), 3);
70-
assert_eq!(uf.find(4), 4);
71-
assert_eq!(uf.find(5), 5);
72-
assert_eq!(uf.find(6), 6);
73-
assert_eq!(uf.find(7), 7);
74-
assert_eq!(uf.find(8), 8);
75-
assert_eq!(uf.find(9), 9);
76-
77-
assert!(uf.union(0, 1));
78-
assert!(uf.union(1, 2));
79-
assert!(uf.union(2, 3));
80-
assert!(uf.union(3, 4));
81-
assert!(uf.union(4, 5));
82-
assert!(uf.union(5, 6));
83-
assert!(uf.union(6, 7));
84-
assert!(uf.union(7, 8));
85-
assert!(uf.union(8, 9));
86-
assert!(!uf.union(9, 0));
177+
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));
192+
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));
87199

88200
assert_eq!(1, uf.count());
89201
}
202+
203+
#[test]
204+
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}
212+
let mut uf = UnionFind::from_iter(["A", "B", "C", "D", "E", "F", "G"]);
213+
uf.union(&"A", &"B");
214+
uf.union(&"B", &"C");
215+
uf.union(&"A", &"D");
216+
uf.union(&"F", &"G");
217+
assert_eq!(3, uf.count());
218+
}
90219
}

0 commit comments

Comments
 (0)