petgraph/
unionfind.rs

1//! `UnionFind<K>` is a disjoint-set data structure.
2
3use 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)]
15pub 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]
30unsafe fn get_unchecked<K>(xs: &[K], index: usize) -> &K
31{
32    debug_assert!(index < xs.len());
33    xs.get_unchecked(index)
34}
35
36impl<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}
petgraph/unionfind.rs

petgraph/
unionfind.rs
