lib/core/collection/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 import hash_collection
  15
  16 # Data structure to keeps track of elements partitioned into disjoint subsets
  17 #
  18 #     var s = new DisjointSet[Int]
  19 #     s.add(1)
  20 #     s.add(2)
  21 #     assert not s.in_same_subset(1,2)
  22 #     s.union(1,2)
  23 #     assert s.in_same_subset(1,2)
  24 #
  25 # `in_same_subset` is transitive, reflexive and symetric
  26 #
  27 #     s.add(3)
  28 #     assert not s.in_same_subset(1,3)
  29 #     s.union(3,2)
  30 #     assert s.in_same_subset(1,3)
  31 #
  32 # Unkike theorical Disjoint-set data structures, the underling implementation is opaque
  33 # that makes the traditionnal `find` method unavailable for clients.
  34 # The methods `in_same_subset`, `to_partitions`, and their variations are offered instead.
  35 class DisjointSet[E]
  36         super SimpleCollection[E]
  37         super Cloneable
  38
  39         # The node in the hiearchical structure for each element
  40         private var nodes = new HashMap[E, DisjointSetNode]
  41
  42         # Copy constructor
  43         init from(other: DisjointSet[E])
  44         do
  45                 # Associate a root node in other to the associated root node in self
  46                 var map = new HashMap[DisjointSetNode, DisjointSetNode]
  47                 for e, v in other.nodes do
  48                         # Create the associated node
  49                         var n2 = new DisjointSetNode
  50                         nodes[e] = n2
  51
  52                         # Get the root node in other and the associated one in self
  53                         var p = other.find(e)
  54                         var p2 = map.get_or_null(p)
  55                         if p2 == null then
  56                                 # if no associated root node, then a new subset is created
  57                                 map[p] = n2.parent
  58                                 number_of_subsets += 1
  59                         else
  60                                 # else attach the new node to the subset of the root node
  61                                 n2.parent = p2
  62                         end
  63                 end
  64         end
  65
  66         # Shallow copy
  67         #
  68         #     var s = new DisjointSet[Int]
  69         #     s.add_all([1,2,3,4,5])
  70         #     s.union_all([1,4,5])
  71         #     var s2 = s.clone
  72         #     assert s2.number_of_subsets == 3
  73         #     assert s2.all_in_same_subset([1,4,5]) == true
  74         #     assert s2.in_same_subset(1,2) == false
  75         redef fun clone do return new DisjointSet[E].from(self)
  76
  77         # The number of subsets in the partition
  78         #
  79         #     var s = new DisjointSet[Int]
  80         #     s.add_all([1,2,3,4,5])
  81         #     assert s.number_of_subsets == 5
  82         #     s.union_all([1,4,5])
  83         #     assert s.number_of_subsets == 3
  84         #     s.union(4,5)
  85         #     assert s.number_of_subsets == 3
  86         var number_of_subsets: Int = 0
  87
  88         # Get the root node of an element
  89         # require: `has(e)`
  90         private fun find(e:E): DisjointSetNode
  91         do
  92                 assert nodes.has_key(e)
  93                 var ne = nodes[e]
  94                 if ne.parent == ne then return ne
  95                 var res = nfind(ne)
  96                 nodes[e] = res
  97                 return res
  98         end
  99
 100         # Get the root node of a node
 101         # Use *path compression* to flatten the structure
 102         # ENSURE: `result.parent == result`
 103         private fun nfind(ne: DisjointSetNode): DisjointSetNode
 104         do
 105                 var nf = ne.parent
 106                 if nf == ne then return ne
 107                 var ng = nfind(nf)
 108                 ne.parent = ng
 109                 return ng
 110         end
 111
 112         # Is the element in the structure
 113         #
 114         #     var s = new DisjointSet[Int]
 115         #     assert not s.has(1)
 116         #     s.add(1)
 117         #     assert s.has(1)
 118         #     assert not s.has(2)
 119         redef fun has(e) do
 120                 return nodes.has_key(e)
 121         end
 122
 123         redef fun iterator do return nodes.keys.iterator
 124
 125         # Add a new element in the structure.
 126         # Initially it is in its own disjoint subset
 127         #
 128         # ENSURE: `has(e)`
 129         redef fun add(e) do
 130                 if nodes.has_key(e) then return
 131                 var ne = new DisjointSetNode
 132                 nodes[e] = ne
 133                 number_of_subsets += 1
 134         end
 135
 136         # Are two elements in the same subset?
 137         fun in_same_subset(e,f:E): Bool
 138         do
 139                 return e == f or find(e) == find(f)
 140         end
 141
 142         # Are all elements of `es` in the same subset?
 143         #
 144         #     var s = new DisjointSet[Int]
 145         #     s.add_all([1,2,3,4,5,6])
 146         #     s.union_all([1,2,3])
 147         #     assert not s.all_in_same_subset([2,3,4])
 148         #     s.union_all([1,4,5])
 149         #     assert s.all_in_same_subset([2,3,4])
 150         fun all_in_same_subset(es: Collection[E]): Bool
 151         do
 152                 if es.is_empty then return true
 153                 var nf = find(es.first)
 154                 for e in es do
 155                         var ne = find(e)
 156                         if ne != nf then return false
 157                 end
 158                 return true
 159         end
 160
 161         # Construct the current partitionning
 162         #
 163         #     var s = new DisjointSet[Int]
 164         #     s.add_all([1,2,3,4,5,6])
 165         #     s.union(1,2)
 166         #     s.union(1,3)
 167         #     s.union(4,5)
 168         #     var p = s.to_partitions
 169         #     assert p.length == 3
 170         fun to_partitions: Collection[Set[E]]
 171         do
 172                 return to_subpartition(self)
 173         end
 174
 175         # Construct a partitionning on `es`, a subset of elements
 176         #
 177         #     var s = new DisjointSet[Int]
 178         #     s.add_all([1,2,3,4,5,6])
 179         #     s.union(1,2)
 180         #     s.union(1,3)
 181         #     s.union(4,5)
 182         #     var p = s.to_subpartition([1,2,4])
 183         #     assert p.length == 2
 184         fun to_subpartition(es: Collection[E]): Collection[Set[E]]
 185         do
 186                 var map = new HashMap[DisjointSetNode, Set[E]]
 187                 for e in es do
 188                         var ne = find(e)
 189                         var set = map.get_or_null(ne)
 190                         if set == null then
 191                                 set = new HashSet[E]
 192                                 map[ne] = set
 193                         end
 194                         set.add(e)
 195                 end
 196                 return map.values
 197         end
 198
 199         # Combine the subsets of `e` and `f`
 200         # ENSURE: `in_same_subset(e,f)`
 201         fun union(e,f:E)
 202         do
 203                 var ne = find(e)
 204                 var nf = find(f)
 205                 if ne == nf then return
 206
 207                 # merge them using *union by rank*
 208                 # attach the smaller tree to the root of the deeper tree
 209                 var er = ne.rank
 210                 var fr = nf.rank
 211                 if er < fr then
 212                         ne.parent = nf
 213                         nodes[e] = nf
 214                 else
 215                         nf.parent = ne
 216                         nodes[f] = ne
 217                         if er == fr then
 218                                 # The only case where the deep is increased is when both are equals
 219                                 ne.rank = er+1
 220                         end
 221                 end
 222                 number_of_subsets -= 1
 223         end
 224
 225         # Combine the subsets of all elements of `es`
 226         # ENSURE: `all_in_same_subset(cs)`
 227         fun union_all(es:Collection[E])
 228         do
 229                 if es.is_empty then return
 230                 var f = es.first
 231                 for e in es do union(e,f)
 232         end
 233 end
 234
 235 # A node in the hierarchical representation of subsets
 236 private class DisjointSetNode
 237         # If parent == self then the node is a root
 238         var parent: DisjointSetNode = self
 239
 240         # The rank to no desequilibrate the structure.
 241         # The term rank is used instead of depth since
 242         # path compression is used, see `DisjointSet::nfind`
 243         var rank = 0
 244 end