lib/union_find.nit

   1 # This file is part of NIT ( http://www.nitlanguage.org ).
   2 #
   3 # This file is free software, which comes along with NIT.  This software is
   4 # distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;
   5 # without  even  the implied warranty of  MERCHANTABILITY or  FITNESS FOR A
   6 # PARTICULAR PURPOSE.  You can modify it is you want,  provided this header
   7 # is kept unaltered, and a notification of the changes is added.
   8 # You  are  allowed  to  redistribute it and sell it, alone or is a part of
   9 # another product.
  10
  11 # union–find algorithm using an efficient disjoint-set data structure
  12 module union_find
  13
  14 # Data structure to keeps track of elements partitioned into disjoint subsets
  15 #     var s = new DisjointSet[Int]
  16 #     s.add(1)
  17 #     s.add(2)
  18 #     assert not s.in_same_subset(1,2)
  19 #     s.union(1,2)
  20 #     assert s.in_same_subset(1,2)
  21 #
  22 # `in_same_subset` is transitive, reflexive and symetric
  23 #
  24 #     s.add(3)
  25 #     assert not s.in_same_subset(1,3)
  26 #     s.union(3,2)
  27 #     assert s.in_same_subset(1,3)
  28 #
  29 # Unkike theorical Disjoint-set data structures, the underling implementation is opaque
  30 # that makes the traditionnal `find` method unavailable for clients.
  31 # The methods `in_same_subset`, `to_partitions`, and their variations are offered instead.
  32 class DisjointSet[E: Object]
  33         super SimpleCollection[E]
  34
  35         # The node in the hiearchical structure for each element
  36         private var nodes = new HashMap[E, DisjointSetNode]
  37
  38         # Get the root node of an element
  39         # require: `has(e)`
  40         private fun find(e:E): DisjointSetNode
  41         do
  42                 assert nodes.has_key(e)
  43                 var ne = nodes[e]
  44                 if ne.parent == ne then return ne
  45                 var res = nfind(ne)
  46                 nodes[e] = res
  47                 return res
  48         end
  49
  50         # Get the root node of a node
  51         # Use *path compression* to flatten the structure
  52         # ENSURE: `result.parent == result`
  53         private fun nfind(ne: DisjointSetNode): DisjointSetNode
  54         do
  55                 var nf = ne.parent
  56                 if nf == ne then return ne
  57                 var ng = nfind(nf)
  58                 ne.parent = ng
  59                 return ng
  60         end
  61
  62         # Is the element in the structure
  63         #
  64         #     var s = new DisjointSet[Int]
  65         #     assert not s.has(1)
  66         #     s.add(1)
  67         #     assert s.has(1)
  68         #     assert not s.has(2)
  69         redef fun has(e: E): Bool
  70         do
  71                 return nodes.has_key(e)
  72         end
  73
  74         redef fun iterator do return nodes.keys.iterator
  75
  76         # Add a new element in the structure.
  77         # Initially it is in its own disjoint subset
  78         #
  79         # ENSURE: `has(e)`
  80         redef fun add(e:E)
  81         do
  82                 if nodes.has_key(e) then return
  83                 var ne = new DisjointSetNode
  84                 nodes[e] = ne
  85         end
  86
  87         # Are two elements in the same subset?
  88         fun in_same_subset(e,f:E): Bool
  89         do
  90                 return e == f or find(e) == find(f)
  91         end
  92
  93         # Are all elements of `es` in the same subset?
  94         #     var s = new DisjointSet[Int]
  95         #     s.add_all([1,2,3,4,5,6])
  96         #     s.union_all([1,2,3])
  97         #     assert not s.all_in_same_subset([2,3,4])
  98         #     s.union_all([1,4,5])
  99         #     assert s.all_in_same_subset([2,3,4])
 100         fun all_in_same_subset(es: Collection[E]): Bool
 101         do
 102                 if es.is_empty then return true
 103                 var nf = find(es.first)
 104                 for e in es do
 105                         var ne = find(e)
 106                         if ne != nf then return false
 107                 end
 108                 return true
 109         end
 110
 111         # Construct the current partitionning
 112         #
 113         #     var s = new DisjointSet[Int]
 114         #     s.add_all([1,2,3,4,5,6])
 115         #     s.union(1,2)
 116         #     s.union(1,3)
 117         #     s.union(4,5)
 118         #     var p = s.to_partitions
 119         #     assert p.length == 3
 120         fun to_partitions: Collection[Set[E]]
 121         do
 122                 return to_subpartition(self)
 123         end
 124
 125         # Construct a partitionning on `es`, a subset of elements
 126         #
 127         #     var s = new DisjointSet[Int]
 128         #     s.add_all([1,2,3,4,5,6])
 129         #     s.union(1,2)
 130         #     s.union(1,3)
 131         #     s.union(4,5)
 132         #     var p = s.to_subpartition([1,2,4])
 133         #     assert p.length == 2
 134         fun to_subpartition(es: Collection[E]): Collection[Set[E]]
 135         do
 136                 var map = new HashMap[DisjointSetNode, Set[E]]
 137                 for e in es do
 138                         var ne = find(e)
 139                         var set = map.get_or_null(ne)
 140                         if set == null then
 141                                 set = new HashSet[E]
 142                                 map[ne] = set
 143                         end
 144                         set.add(e)
 145                 end
 146                 return map.values
 147         end
 148
 149         # Combine the subsets of `e` and `f`
 150         # ENSURE: `in_same_subset(e,f)`
 151         fun union(e,f:E)
 152         do
 153                 var ne = find(e)
 154                 var nf = find(f)
 155                 if ne == nf then return
 156
 157                 # merge them using *union by rank*
 158                 # attach the smaller tree to the root of the deeper tree
 159                 var er = ne.rank
 160                 var fr = nf.rank
 161                 if er < fr then
 162                         ne.parent = nf
 163                         nodes[e] = nf
 164                 else
 165                         nf.parent = ne
 166                         nodes[f] = ne
 167                         if er == fr then
 168                                 # The only case where the deep is increased is when both are equals
 169                                 ne.rank = er+1
 170                         end
 171                 end
 172         end
 173
 174         # Combine the subsets of all elements of `es`
 175         # ENSURE: `all_in_same_subset(cs)`
 176         fun union_all(es:Collection[E])
 177         do
 178                 if es.is_empty then return
 179                 var f = es.first
 180                 for e in es do union(e,f)
 181         end
 182 end
 183
 184 # A node in the hierarchical representation of subsets
 185 private class DisjointSetNode
 186         # If parent == self then the node is a root
 187         var parent: DisjointSetNode = self
 188
 189         # The rank to no desequilibrate the structure.
 190         # The term rank is used instead of depth since
 191         # path compression is used, see `DisjointSet::nfind`
 192         var rank = 0
 193 end