1 //! `UnionFind<K>` is a disjoint-set data structure.
2
3 use super::graph::IndexType;
4 use std::cmp::Ordering;
5
6 /// `UnionFind<K>` is a disjoint-set data structure. It tracks set membership of *n* elements
7 /// indexed from *0* to *n - 1*. The scalar type is `K` which must be an unsigned integer type.
8 ///
9 /// <http://en.wikipedia.org/wiki/Disjoint-set_data_structure>
10 ///
11 /// Too awesome not to quote:
12 ///
13 /// “The amortized time per operation is **O(α(n))** where **α(n)** is the
14 /// inverse of **f(x) = A(x, x)** with **A** being the extremely fast-growing Ackermann function.”
15 #[derive(Debug, Clone)]
16 pub struct UnionFind<K> {
17 // For element at index *i*, store the index of its parent; the representative itself
18 // stores its own index. This forms equivalence classes which are the disjoint sets, each
19 // with a unique representative.
20 parent: Vec<K>,
21 // It is a balancing tree structure,
22 // so the ranks are logarithmic in the size of the container -- a byte is more than enough.
23 //
24 // Rank is separated out both to save space and to save cache in when searching in the parent
25 // vector.
26 rank: Vec<u8>,
27 }
28
29 #[inline]
get_unchecked<K>(xs: &[K], index: usize) -> &K30 unsafe fn get_unchecked<K>(xs: &[K], index: usize) -> &K {
31 debug_assert!(index < xs.len());
32 xs.get_unchecked(index)
33 }
34
35 #[inline]
get_unchecked_mut<K>(xs: &mut [K], index: usize) -> &mut K36 unsafe fn get_unchecked_mut<K>(xs: &mut [K], index: usize) -> &mut K {
37 debug_assert!(index < xs.len());
38 xs.get_unchecked_mut(index)
39 }
40
41 impl<K> UnionFind<K>
42 where
43 K: IndexType,
44 {
45 /// Create a new `UnionFind` of `n` disjoint sets.
new(n: usize) -> Self46 pub fn new(n: usize) -> Self {
47 let rank = vec![0; n];
48 let parent = (0..n).map(K::new).collect::<Vec<K>>();
49
50 UnionFind { parent, rank }
51 }
52
53 /// Return the representative for `x`.
54 ///
55 /// **Panics** if `x` is out of bounds.
find(&self, x: K) -> K56 pub fn find(&self, x: K) -> K {
57 assert!(x.index() < self.parent.len());
58 unsafe {
59 let mut x = x;
60 loop {
61 // Use unchecked indexing because we can trust the internal set ids.
62 let xparent = *get_unchecked(&self.parent, x.index());
63 if xparent == x {
64 break;
65 }
66 x = xparent;
67 }
68 x
69 }
70 }
71
72 /// Return the representative for `x`.
73 ///
74 /// Write back the found representative, flattening the internal
75 /// datastructure in the process and quicken future lookups.
76 ///
77 /// **Panics** if `x` is out of bounds.
find_mut(&mut self, x: K) -> K78 pub fn find_mut(&mut self, x: K) -> K {
79 assert!(x.index() < self.parent.len());
80 unsafe { self.find_mut_recursive(x) }
81 }
82
find_mut_recursive(&mut self, mut x: K) -> K83 unsafe fn find_mut_recursive(&mut self, mut x: K) -> K {
84 let mut parent = *get_unchecked(&self.parent, x.index());
85 while parent != x {
86 let grandparent = *get_unchecked(&self.parent, parent.index());
87 *get_unchecked_mut(&mut self.parent, x.index()) = grandparent;
88 x = parent;
89 parent = grandparent;
90 }
91 x
92 }
93
94 /// Returns `true` if the given elements belong to the same set, and returns
95 /// `false` otherwise.
equiv(&self, x: K, y: K) -> bool96 pub fn equiv(&self, x: K, y: K) -> bool {
97 self.find(x) == self.find(y)
98 }
99
100 /// Unify the two sets containing `x` and `y`.
101 ///
102 /// Return `false` if the sets were already the same, `true` if they were unified.
103 ///
104 /// **Panics** if `x` or `y` is out of bounds.
union(&mut self, x: K, y: K) -> bool105 pub fn union(&mut self, x: K, y: K) -> bool {
106 if x == y {
107 return false;
108 }
109 let xrep = self.find_mut(x);
110 let yrep = self.find_mut(y);
111
112 if xrep == yrep {
113 return false;
114 }
115
116 let xrepu = xrep.index();
117 let yrepu = yrep.index();
118 let xrank = self.rank[xrepu];
119 let yrank = self.rank[yrepu];
120
121 // The rank corresponds roughly to the depth of the treeset, so put the
122 // smaller set below the larger
123 match xrank.cmp(&yrank) {
124 Ordering::Less => self.parent[xrepu] = yrep,
125 Ordering::Greater => self.parent[yrepu] = xrep,
126 Ordering::Equal => {
127 self.parent[yrepu] = xrep;
128 self.rank[xrepu] += 1;
129 }
130 }
131 true
132 }
133
134 /// Return a vector mapping each element to its representative.
into_labeling(mut self) -> Vec<K>135 pub fn into_labeling(mut self) -> Vec<K> {
136 // write in the labeling of each element
137 unsafe {
138 for ix in 0..self.parent.len() {
139 let k = *get_unchecked(&self.parent, ix);
140 let xrep = self.find_mut_recursive(k);
141 *self.parent.get_unchecked_mut(ix) = xrep;
142 }
143 }
144 self.parent
145 }
146 }
147