lib/standard/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 #     var s = new DisjointSet[Int]
  18 #     s.add(1)
  19 #     s.add(2)
  20 #     assert not s.in_same_subset(1,2)
  21 #     s.union(1,2)
  22 #     assert s.in_same_subset(1,2)
  23 #
  24 # `in_same_subset` is transitive, reflexive and symetric
  25 #
  26 #     s.add(3)
  27 #     assert not s.in_same_subset(1,3)
  28 #     s.union(3,2)
  29 #     assert s.in_same_subset(1,3)
  30 #
  31 # Unkike theorical Disjoint-set data structures, the underling implementation is opaque
  32 # that makes the traditionnal `find` method unavailable for clients.
  33 # The methods `in_same_subset`, `to_partitions`, and their variations are offered instead.
  34 class DisjointSet[E]
  35         super SimpleCollection[E]
  36         super Cloneable
  37
  38         # The node in the hiearchical structure for each element
  39         private var nodes = new HashMap[E, DisjointSetNode]
  40
  41         # Copy constructor
  42         init from(other: DisjointSet[E])
  43         do
  44                 # Associate a root node in other to the associated root node in self
  45                 var map = new HashMap[DisjointSetNode, DisjointSetNode]
  46                 for e, v in other.nodes do
  47                         # Create the associated node
  48                         var n2 = new DisjointSetNode
  49                         nodes[e] = n2
  50
  51                         # Get the root node in other and the associated one in self
  52                         var p = other.find(e)
  53                         var p2 = map.get_or_null(p)
  54                         if p2 == null then
  55                                 # if no associated root node, then a new subset is created
  56                                 map[p] = n2.parent
  57                                 number_of_subsets += 1
  58                         else
  59                                 # else attach the new node to the subset of the root node
  60                                 n2.parent = p2
  61                         end
  62                 end
  63         end
  64
  65         # Shallow copy
  66         #
  67         #     var s = new DisjointSet[Int]
  68         #     s.add_all([1,2,3,4,5])
  69         #     s.union_all([1,4,5])
  70         #     var s2 = s.clone
  71         #     assert s2.number_of_subsets == 3
  72         #     assert s2.all_in_same_subset([1,4,5]) == true
  73         #     assert s2.in_same_subset(1,2) == false
  74         redef fun clone do return new DisjointSet[E].from(self)
  75
  76         # The number of subsets in the partition
  77         #
  78         #     var s = new DisjointSet[Int]
  79         #     s.add_all([1,2,3,4,5])
  80         #     assert s.number_of_subsets == 5
  81         #     s.union_all([1,4,5])
  82         #     assert s.number_of_subsets == 3
  83         #     s.union(4,5)
  84         #     assert s.number_of_subsets == 3
  85         var number_of_subsets: Int = 0
  86
  87         # Get the root node of an element
  88         # require: `has(e)`
  89         private fun find(e:E): DisjointSetNode
  90         do
  91                 assert nodes.has_key(e)
  92                 var ne = nodes[e]
  93                 if ne.parent == ne then return ne
  94                 var res = nfind(ne)
  95                 nodes[e] = res
  96                 return res
  97         end
  98
  99         # Get the root node of a node
 100         # Use *path compression* to flatten the structure
 101         # ENSURE: `result.parent == result`
 102         private fun nfind(ne: DisjointSetNode): DisjointSetNode
 103         do
 104                 var nf = ne.parent
 105                 if nf == ne then return ne
 106                 var ng = nfind(nf)
 107                 ne.parent = ng
 108                 return ng
 109         end
 110
 111         # Is the element in the structure
 112         #
 113         #     var s = new DisjointSet[Int]
 114         #     assert not s.has(1)
 115         #     s.add(1)
 116         #     assert s.has(1)
 117         #     assert not s.has(2)
 118         redef fun has(e: E): Bool
 119         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:E)
 130         do
 131                 if nodes.has_key(e) then return
 132                 var ne = new DisjointSetNode
 133                 nodes[e] = ne
 134                 number_of_subsets += 1
 135         end
 136
 137         # Are two elements in the same subset?
 138         fun in_same_subset(e,f:E): Bool
 139         do
 140                 return e == f or find(e) == find(f)
 141         end
 142
 143         # Are all elements of `es` in the same subset?
 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