yui_core/algo/
union_find.rs

1//! Disjoint-set structures: [`UnionFind`] over `0..n` (path compression and
2//! union by rank, from `petgraph`), and [`KeyedUnionFind`] over hashable keys.
3
4use std::collections::HashMap;
5use std::hash::Hash;
6
7use indexmap::IndexSet;
8use petgraph::unionfind::UnionFind as PgUnionFind;
9
10/// A disjoint-set data structure over `0..n`. Backed by [`petgraph::unionfind::UnionFind`],
11/// which uses path compression and union by rank for `O(α(n))` amortized operations.
12#[derive(Clone)]
13pub struct UnionFind {
14    inner: PgUnionFind<usize>,
15}
16
17impl UnionFind {
18    pub fn new(n: usize) -> Self {
19        Self { inner: PgUnionFind::new(n) }
20    }
21
22    pub fn extend(&mut self, l: usize) {
23        for _ in 0..l {
24            self.inner.new_set();
25        }
26    }
27
28    pub fn size(&self) -> usize {
29        self.inner.len()
30    }
31
32    pub fn root(&self, i: usize) -> usize {
33        self.inner.find(i)
34    }
35
36    pub fn is_same(&self, i: usize, j: usize) -> bool {
37        self.inner.equiv(i, j)
38    }
39
40    pub fn union(&mut self, i: usize, j: usize) {
41        self.inner.union(i, j);
42    }
43
44    pub fn into_disjoint(self) -> Vec<Vec<usize>> {
45        let labeling = self.inner.into_labeling();
46        let n = labeling.len();
47        labeling.iter().enumerate().fold(
48            (Vec::<Vec<usize>>::new(), vec![None::<usize>; n]),
49            |(mut groups, mut idx_of), (i, &r)| {
50                let idx = *idx_of[r].get_or_insert_with(|| {
51                    groups.push(vec![]);
52                    groups.len() - 1
53                });
54                groups[idx].push(i);
55                (groups, idx_of)
56            },
57        ).0
58    }
59}
60
61pub struct KeyedUnionFind<X> where X: Eq + Hash {
62    inner: UnionFind,
63    keys: IndexSet<X>,
64}
65
66impl<X> KeyedUnionFind<X> where X: Eq + Hash {
67    pub fn new() -> Self {
68        Self { inner: UnionFind::new(0), keys: IndexSet::new() }
69    }
70
71    /// Insert a key, returning its index. Returns the existing index if `x` is already present.
72    pub fn insert(&mut self, x: X) -> usize {
73        let (idx, was_new) = self.keys.insert_full(x);
74        if was_new {
75            self.inner.extend(1);
76        }
77        idx
78    }
79
80    fn index_of(&self, x: &X) -> usize {
81        self.keys.get_index_of(x).expect("key is not in the union-find")
82    }
83
84    fn element_at(&self, i: usize) -> &X {
85        &self.keys[i]
86    }
87
88    pub fn size(&self) -> usize {
89        self.inner.size()
90    }
91
92    pub fn contains(&self, x: &X) -> bool {
93        self.keys.contains(x)
94    }
95
96    pub fn root(&self, x: &X) -> &X {
97        let i = self.index_of(x);
98        let j = self.inner.root(i);
99        self.element_at(j)
100    }
101
102    pub fn is_same(&self, x: &X, y: &X) -> bool {
103        self.root(x) == self.root(y)
104    }
105
106    pub fn union(&mut self, x: &X, y: &X) {
107        let i = self.index_of(x);
108        let j = self.index_of(y);
109        self.inner.union(i, j);
110    }
111
112    pub fn into_disjoint(self) -> Vec<Vec<X>> {
113        let Self { inner, keys } = self;
114        let group = inner.into_disjoint();
115
116        let mut map: HashMap<usize, X> = keys.into_iter().enumerate().collect();
117
118        group.iter().map(|l|
119            l.iter().map(|i| map.remove(i).unwrap()).collect()
120        ).collect()
121    }
122}
123
124impl<X> FromIterator<X> for KeyedUnionFind<X>
125where X: Hash + Eq {
126    fn from_iter<T: IntoIterator<Item = X>>(keys: T) -> Self {
127        let keys: IndexSet<X> = keys.into_iter().collect();
128        let inner = UnionFind::new(keys.len());
129        Self { inner, keys }
130    }
131}
132
133#[cfg(test)]
134mod tests {
135    use std::collections::HashSet;
136    use super::*;
137
138    fn group_as_sets(g: Vec<Vec<usize>>) -> HashSet<Vec<usize>> {
139        g.into_iter().map(|mut v| { v.sort(); v }).collect()
140    }
141
142    #[test]
143    fn test() {
144        let mut u = UnionFind::new(4);
145
146        assert_eq!(u.size(), 4);
147        assert!(!u.is_same(0, 1));
148        assert!(!u.is_same(1, 2));
149        assert!(!u.is_same(2, 3));
150        assert_eq!(group_as_sets(u.clone().into_disjoint()), HashSet::from([vec![0], vec![1], vec![2], vec![3]]));
151
152        u.union(0, 1);
153
154        assert!( u.is_same(0, 1));
155        assert!(!u.is_same(1, 2));
156        assert!(!u.is_same(2, 3));
157        assert_eq!(group_as_sets(u.clone().into_disjoint()), HashSet::from([vec![0, 1], vec![2], vec![3]]));
158
159        u.union(2, 3);
160
161        assert!( u.is_same(0, 1));
162        assert!(!u.is_same(1, 2));
163        assert!( u.is_same(2, 3));
164        assert_eq!(group_as_sets(u.clone().into_disjoint()), HashSet::from([vec![0, 1], vec![2, 3]]));
165
166        u.union(1, 3);
167
168        assert!( u.is_same(0, 1));
169        assert!( u.is_same(1, 2));
170        assert!( u.is_same(2, 3));
171        assert_eq!(group_as_sets(u.clone().into_disjoint()), HashSet::from([vec![0, 1, 2, 3]]));
172    }
173
174    fn keyed_with_unions(unions: &[(&'static str, &'static str)]) -> KeyedUnionFind<&'static str> {
175        let mut u = KeyedUnionFind::from_iter(["a", "b", "c", "d"]);
176        for (x, y) in unions { u.union(x, y); }
177        u
178    }
179
180    fn sorted_disjoint(u: KeyedUnionFind<&'static str>) -> HashSet<Vec<&'static str>> {
181        u.into_disjoint().into_iter().map(|mut v| { v.sort(); v }).collect()
182    }
183
184    #[test]
185    fn test_hash_no_unions() {
186        let u = keyed_with_unions(&[]);
187        assert_eq!(u.size(), 4);
188        assert!(!u.is_same(&"a", &"b"));
189        assert!(!u.is_same(&"b", &"c"));
190        assert!(!u.is_same(&"c", &"d"));
191        assert_eq!(sorted_disjoint(u), HashSet::from([vec!["a"], vec!["b"], vec!["c"], vec!["d"]]));
192    }
193
194    #[test]
195    fn test_hash_one_union() {
196        let u = keyed_with_unions(&[("a", "b")]);
197        assert!( u.is_same(&"a", &"b"));
198        assert!(!u.is_same(&"b", &"c"));
199        assert!(!u.is_same(&"c", &"d"));
200        assert_eq!(sorted_disjoint(u), HashSet::from([vec!["a", "b"], vec!["c"], vec!["d"]]));
201    }
202
203    #[test]
204    fn test_hash_two_unions() {
205        let u = keyed_with_unions(&[("a", "b"), ("c", "d")]);
206        assert!( u.is_same(&"a", &"b"));
207        assert!(!u.is_same(&"b", &"c"));
208        assert!( u.is_same(&"c", &"d"));
209        assert_eq!(sorted_disjoint(u), HashSet::from([vec!["a", "b"], vec!["c", "d"]]));
210    }
211
212    #[test]
213    fn test_hash_all_unioned() {
214        let u = keyed_with_unions(&[("a", "b"), ("c", "d"), ("b", "d")]);
215        assert!(u.is_same(&"a", &"b"));
216        assert!(u.is_same(&"b", &"c"));
217        assert!(u.is_same(&"c", &"d"));
218        assert_eq!(sorted_disjoint(u), HashSet::from([vec!["a", "b", "c", "d"]]));
219    }
220}
yui_core/algo/union_find.rs

yui_core/algo/
union_find.rs
