]>
Commit | Line | Data |
---|---|---|
83c7162d XL |
1 | //! `UnionFind<K>` is a disjoint-set data structure. |
2 | ||
3 | use super::graph::IndexType; | |
4 | ||
5 | /// `UnionFind<K>` is a disjoint-set data structure. It tracks set membership of *n* elements | |
6 | /// indexed from *0* to *n - 1*. The scalar type is `K` which must be an unsigned integer type. | |
7 | /// | |
8 | /// <http://en.wikipedia.org/wiki/Disjoint-set_data_structure> | |
9 | /// | |
10 | /// Too awesome not to quote: | |
11 | /// | |
12 | /// “The amortized time per operation is **O(α(n))** where **α(n)** is the | |
13 | /// inverse of **f(x) = A(x, x)** with **A** being the extremely fast-growing Ackermann function.” | |
14 | #[derive(Debug, Clone)] | |
15 | pub struct UnionFind<K> | |
16 | { | |
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] | |
30 | unsafe fn get_unchecked<K>(xs: &[K], index: usize) -> &K | |
31 | { | |
32 | debug_assert!(index < xs.len()); | |
33 | xs.get_unchecked(index) | |
34 | } | |
35 | ||
36 | impl<K> UnionFind<K> | |
37 | where K: IndexType | |
38 | { | |
39 | /// Create a new `UnionFind` of `n` disjoint sets. | |
40 | pub fn new(n: usize) -> Self | |
41 | { | |
42 | let rank = vec![0; n]; | |
43 | let parent = (0..n).map(K::new).collect::<Vec<K>>(); | |
44 | ||
45 | UnionFind{parent: parent, rank: rank} | |
46 | } | |
47 | ||
48 | /// Return the representative for `x`. | |
49 | /// | |
50 | /// **Panics** if `x` is out of bounds. | |
51 | pub fn find(&self, x: K) -> K | |
52 | { | |
53 | assert!(x.index() < self.parent.len()); | |
54 | unsafe { | |
55 | let mut x = x; | |
56 | loop { | |
57 | // Use unchecked indexing because we can trust the internal set ids. | |
58 | let xparent = *get_unchecked(&self.parent, x.index()); | |
59 | if xparent == x { | |
60 | break | |
61 | } | |
62 | x = xparent; | |
63 | } | |
64 | x | |
65 | } | |
66 | } | |
67 | ||
68 | /// Return the representative for `x`. | |
69 | /// | |
70 | /// Write back the found representative, flattening the internal | |
71 | /// datastructure in the process and quicken future lookups. | |
72 | /// | |
73 | /// **Panics** if `x` is out of bounds. | |
74 | pub fn find_mut(&mut self, x: K) -> K | |
75 | { | |
76 | assert!(x.index() < self.parent.len()); | |
77 | unsafe { | |
78 | self.find_mut_recursive(x) | |
79 | } | |
80 | } | |
81 | ||
82 | unsafe fn find_mut_recursive(&mut self, x: K) -> K | |
83 | { | |
84 | let xparent = *get_unchecked(&self.parent, x.index()); | |
85 | if xparent != x { | |
86 | let xrep = self.find_mut_recursive(xparent); | |
87 | let xparent = self.parent.get_unchecked_mut(x.index()); | |
88 | *xparent = xrep; | |
89 | *xparent | |
90 | } else { | |
91 | xparent | |
92 | } | |
93 | } | |
94 | ||
95 | ||
96 | /// Unify the two sets containing `x` and `y`. | |
97 | /// | |
98 | /// Return `false` if the sets were already the same, `true` if they were unified. | |
99 | /// | |
100 | /// **Panics** if `x` or `y` is out of bounds. | |
101 | pub fn union(&mut self, x: K, y: K) -> bool | |
102 | { | |
103 | if x == y { | |
104 | return false | |
105 | } | |
106 | let xrep = self.find_mut(x); | |
107 | let yrep = self.find_mut(y); | |
108 | ||
109 | if xrep == yrep { | |
110 | return false | |
111 | } | |
112 | ||
113 | let xrepu = xrep.index(); | |
114 | let yrepu = yrep.index(); | |
115 | let xrank = self.rank[xrepu]; | |
116 | let yrank = self.rank[yrepu]; | |
117 | ||
118 | // The rank corresponds roughly to the depth of the treeset, so put the | |
119 | // smaller set below the larger | |
120 | if xrank < yrank { | |
121 | self.parent[xrepu] = yrep; | |
122 | } else if xrank > yrank { | |
123 | self.parent[yrepu] = xrep; | |
124 | } else { | |
125 | // put y below x when equal. | |
126 | self.parent[yrepu] = xrep; | |
127 | self.rank[xrepu] += 1; | |
128 | } | |
129 | true | |
130 | } | |
131 | ||
132 | /// Return a vector mapping each element to its representative. | |
133 | pub fn into_labeling(mut self) -> Vec<K> | |
134 | { | |
135 | // write in the labeling of each element | |
136 | unsafe { | |
137 | for ix in 0..self.parent.len() { | |
138 | let k = *get_unchecked(&self.parent, ix); | |
139 | let xrep = self.find_mut_recursive(k); | |
140 | *self.parent.get_unchecked_mut(ix) = xrep; | |
141 | } | |
142 | } | |
143 | self.parent | |
144 | } | |
145 | } |