petgraph/
unionfind.rs

1//! `UnionFind<K>` is a disjoint-set data structure.
2
3use super::graph::IndexType;
4use 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)]
16pub 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]
30unsafe fn get_unchecked<K>(xs: &[K], index: usize) -> &K {
31    debug_assert!(index < xs.len());
32    xs.get_unchecked(index)
33}
34
35#[inline]
36unsafe 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
41impl<K> UnionFind<K>
42where
43    K: IndexType,
44{
45    /// Create a new `UnionFind` of `n` disjoint sets.
46    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.
56    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.
78    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
83    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.
96    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.
105    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.
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}
petgraph/unionfind.rs

petgraph/
unionfind.rs
