vendor/petgraph/src/unionfind.rs

   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 }