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 keep 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 symmetric
  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 # Unlike theoretical Disjoint-set data structures, the underling implementation is opaque
  33 # making the traditional `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         #     var s = new DisjointSet[Int]
 144         #     s.add_all([1,2,3,4,5,6])
 145         #     s.union_all([1,2,3])
 146         #     assert not s.all_in_same_subset([2,3,4])
 147         #     s.union_all([1,4,5])
 148         #     assert s.all_in_same_subset([2,3,4])
 149         fun all_in_same_subset(es: Collection[E]): Bool
 150         do
 151                 if es.is_empty then return true
 152                 var nf = find(es.first)
 153                 for e in es do
 154                         var ne = find(e)
 155                         if ne != nf then return false
 156                 end
 157                 return true
 158         end
 159
 160         # Construct the current partitionning
 161         #
 162         #     var s = new DisjointSet[Int]
 163         #     s.add_all([1,2,3,4,5,6])
 164         #     s.union(1,2)
 165         #     s.union(1,3)
 166         #     s.union(4,5)
 167         #     var p = s.to_partitions
 168         #     assert p.length == 3
 169         fun to_partitions: Collection[Set[E]]
 170         do
 171                 return to_subpartition(self)
 172         end
 173
 174         # Construct a partitioning on `es`, a subset of elements
 175         #
 176         #     var s = new DisjointSet[Int]
 177         #     s.add_all([1,2,3,4,5,6])
 178         #     s.union(1,2)
 179         #     s.union(1,3)
 180         #     s.union(4,5)
 181         #     var p = s.to_subpartition([1,2,4])
 182         #     assert p.length == 2
 183         fun to_subpartition(es: Collection[E]): Collection[Set[E]]
 184         do
 185                 var map = new HashMap[DisjointSetNode, Set[E]]
 186                 for e in es do
 187                         var ne = find(e)
 188                         var set = map.get_or_null(ne)
 189                         if set == null then
 190                                 set = new HashSet[E]
 191                                 map[ne] = set
 192                         end
 193                         set.add(e)
 194                 end
 195                 return map.values
 196         end
 197
 198         # Combine the subsets of `e` and `f`
 199         # ENSURE: `in_same_subset(e,f)`
 200         fun union(e,f:E)
 201         do
 202                 var ne = find(e)
 203                 var nf = find(f)
 204                 if ne == nf then return
 205
 206                 # merge them using *union by rank*
 207                 # attach the smaller tree to the root of the deeper tree
 208                 var er = ne.rank
 209                 var fr = nf.rank
 210                 if er < fr then
 211                         ne.parent = nf
 212                         nodes[e] = nf
 213                 else
 214                         nf.parent = ne
 215                         nodes[f] = ne
 216                         if er == fr then
 217                                 # The only case where the deep is increased is when both are equals
 218                                 ne.rank = er+1
 219                         end
 220                 end
 221                 number_of_subsets -= 1
 222         end
 223
 224         # Combine the subsets of all elements of `es`
 225         # ENSURE: `all_in_same_subset(cs)`
 226         fun union_all(es:Collection[E])
 227         do
 228                 if es.is_empty then return
 229                 var f = es.first
 230                 for e in es do union(e,f)
 231         end
 232 end
 233
 234 # A node in the hierarchical representation of subsets
 235 private class DisjointSetNode
 236         # If parent == self then the node is a root
 237         var parent: DisjointSetNode = self
 238
 239         # The rank to keep the structure balanced.
 240         # The term rank is used instead of depth since
 241         # path compression is used, see `DisjointSet::nfind`
 242         var rank = 0
 243 end