petgraph/
unionfind.rs

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

petgraph/
unionfind.rs
