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 #     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) do
 119                 return nodes.has_key(e)
 120         end
 121
 122         redef fun iterator do return nodes.keys.iterator
 123
 124         # Add a new element in the structure.
 125         # Initially it is in its own disjoint subset
 126         #
 127         # ENSURE: `has(e)`
 128         redef fun add(e) do
 129                 if nodes.has_key(e) then return
 130                 var ne = new DisjointSetNode
 131                 nodes[e] = ne
 132                 number_of_subsets += 1
 133         end
 134
 135         # Are two elements in the same subset?
 136         fun in_same_subset(e,f:E): Bool
 137         do
 138                 return e == f or find(e) == find(f)
 139         end
 140
 141         # Are all elements of `es` in the same subset?
 142         #     var s = new DisjointSet[Int]
 143         #     s.add_all([1,2,3,4,5,6])
 144         #     s.union_all([1,2,3])
 145         #     assert not s.all_in_same_subset([2,3,4])
 146         #     s.union_all([1,4,5])
 147         #     assert s.all_in_same_subset([2,3,4])
 148         fun all_in_same_subset(es: Collection[E]): Bool
 149         do
 150                 if es.is_empty then return true
 151                 var nf = find(es.first)
 152                 for e in es do
 153                         var ne = find(e)
 154                         if ne != nf then return false
 155                 end
 156                 return true
 157         end
 158
 159         # Construct the current partitionning
 160         #
 161         #     var s = new DisjointSet[Int]
 162         #     s.add_all([1,2,3,4,5,6])
 163         #     s.union(1,2)
 164         #     s.union(1,3)
 165         #     s.union(4,5)
 166         #     var p = s.to_partitions
 167         #     assert p.length == 3
 168         fun to_partitions: Collection[Set[E]]
 169         do
 170                 return to_subpartition(self)
 171         end
 172
 173         # Construct a partitionning on `es`, a subset of elements
 174         #
 175         #     var s = new DisjointSet[Int]
 176         #     s.add_all([1,2,3,4,5,6])
 177         #     s.union(1,2)
 178         #     s.union(1,3)
 179         #     s.union(4,5)
 180         #     var p = s.to_subpartition([1,2,4])
 181         #     assert p.length == 2
 182         fun to_subpartition(es: Collection[E]): Collection[Set[E]]
 183         do
 184                 var map = new HashMap[DisjointSetNode, Set[E]]
 185                 for e in es do
 186                         var ne = find(e)
 187                         var set = map.get_or_null(ne)
 188                         if set == null then
 189                                 set = new HashSet[E]
 190                                 map[ne] = set
 191                         end
 192                         set.add(e)
 193                 end
 194                 return map.values
 195         end
 196
 197         # Combine the subsets of `e` and `f`
 198         # ENSURE: `in_same_subset(e,f)`
 199         fun union(e,f:E)
 200         do
 201                 var ne = find(e)
 202                 var nf = find(f)
 203                 if ne == nf then return
 204
 205                 # merge them using *union by rank*
 206                 # attach the smaller tree to the root of the deeper tree
 207                 var er = ne.rank
 208                 var fr = nf.rank
 209                 if er < fr then
 210                         ne.parent = nf
 211                         nodes[e] = nf
 212                 else
 213                         nf.parent = ne
 214                         nodes[f] = ne
 215                         if er == fr then
 216                                 # The only case where the deep is increased is when both are equals
 217                                 ne.rank = er+1
 218                         end
 219                 end
 220                 number_of_subsets -= 1
 221         end
 222
 223         # Combine the subsets of all elements of `es`
 224         # ENSURE: `all_in_same_subset(cs)`
 225         fun union_all(es:Collection[E])
 226         do
 227                 if es.is_empty then return
 228                 var f = es.first
 229                 for e in es do union(e,f)
 230         end
 231 end
 232
 233 # A node in the hierarchical representation of subsets
 234 private class DisjointSetNode
 235         # If parent == self then the node is a root
 236         var parent: DisjointSetNode = self
 237
 238         # The rank to no desequilibrate the structure.
 239         # The term rank is used instead of depth since
 240         # path compression is used, see `DisjointSet::nfind`
 241         var rank = 0
 242 end