lib/trees/bintree.nit

   1 # This file is part of NIT ( http://www.nitlanguage.org ).
   2 #
   3 # Licensed under the Apache License, Version 2.0 (the "License");
   4 # you may not use this file except in compliance with the License.
   5 # You may obtain a copy of the License at
   6 #
   7 #     http://www.apache.org/licenses/LICENSE-2.0
   8 #
   9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS,
  11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 # See the License for the specific language governing permissions and
  13 # limitations under the License.
  14
  15 # Binary Tree data-structure
  16 # A binary tree is a tree data structure in which each node has at most two children
  17 # (referred to as the left child and the right child).
  18 # In a binary tree, the degree of each node can be at most two.
  19 # Binary trees are used to implement binary search trees and binary heaps,
  20 # and are used for efficient searching and sorting.
  21 module bintree
  22
  23 import abstract_tree
  24
  25 # Binary Tree Map
  26 #
  27 # Properties:
  28 #  * unique root
  29 #  * node.left.key < node.key
  30 #  * node.right.key > node.key
  31 #  * no duplicates allowed
  32 #
  33 # Operations:
  34 #  * search average O(lg n) worst O(n)
  35 #  * insert average O(lg n) worst O(n)
  36 #  * delete average O(lg n) worst O(n)
  37 #
  38 # Usage:
  39 #     var tree = new BinTreeMap[Int, String]
  40 #     tree[1] = "n1"
  41 #     assert tree.min == "n1"
  42 class BinTreeMap[K: Comparable, E]
  43         super TreeMap[K, E]
  44
  45         redef type N: BinTreeNode[K, E]
  46
  47         private var len = 0
  48         private var first_node: nullable BinTreeNode[K, E] = null
  49         private var last_node: nullable BinTreeNode[K, E] = null
  50
  51         # O(n) in worst case, average is O(h) with h: tree height
  52         #
  53         #     var tree = new BinTreeMap[Int, String]
  54         #     assert tree.is_empty
  55         #     tree[1] = "n1"
  56         #     assert not tree.is_empty
  57         redef fun is_empty do return root == null
  58
  59         # O(n) in worst case, average is O(h) with h: tree height
  60         #
  61         #     var tree = new BinTreeMap[Int, String]
  62         #     assert not tree.has_key(1)
  63         #     for i in [4, 2, 1, 5, 3] do tree[i] = "n{i}"
  64         #     assert not tree.has_key(0)
  65         #     assert tree.has_key(2)
  66         #     assert not tree.has_key(6)
  67         redef fun has_key(key: K): Bool do
  68                 if is_empty then return false
  69                 var res = search_down(root.as(not null), key)
  70                 if res != null then
  71                         cache_node = res
  72                         return true
  73                 end
  74                 return false
  75         end
  76
  77         private var cache_node: nullable N = null
  78
  79         # Get the node value associated to `key`
  80         # O(n) in worst case, average is O(h) with h: tree height
  81         #
  82         #     var tree = new BinTreeMap[Int, String]
  83         #     for i in [4, 2, 1, 5, 3] do tree[i] = "n{i}"
  84         #     assert tree.has_key(1)
  85         #     assert tree[1] == "n1"
  86         #     assert tree.has_key(1)
  87         #     assert tree[2] == "n2"
  88         redef fun [](key: K): E do
  89                 assert not_empty: not is_empty
  90                 if cache_node != null and cache_node.key == key then return cache_node.value
  91                 var res = search_down(root.as(not null), key)
  92                 assert has_key: res != null
  93                 return res.value
  94         end
  95
  96         protected fun search_down(from: N, key: K): nullable N do
  97                 if key == from.key then return from
  98                 if from.left != null and key < from.key then
  99                         return search_down(from.left.as(not null), key)
 100                 else if from.right != null then
 101                         return search_down(from.right.as(not null), key)
 102                 end
 103                 return null
 104         end
 105
 106         # Get the node with the minimum key
 107         # O(n) in worst case, average is O(h) with h: tree height
 108         #
 109         #     var tree = new BinTreeMap[Int, String]
 110         #     for i in [4, 2, 1, 5, 3] do tree[i] = "n{i}"
 111         #     assert tree.min == "n1"
 112         fun min: E do
 113                 assert not_empty: root != null
 114                 return min_from(root.as(not null)).value
 115         end
 116
 117         protected fun min_from(node: N): N do
 118                 if node.left == null then return node
 119                 return min_from(node.left.as(not null))
 120         end
 121
 122         # Get the node with the maximum key
 123         # O(n) in worst case, average is O(h) with h: tree height
 124         #
 125         #     var tree = new BinTreeMap[Int, String]
 126         #     for i in [4, 2, 1, 5, 3, 6, 7, 8] do tree[i] = "n{i}"
 127         #     assert tree.max == "n8"
 128         fun max: E do
 129                 assert not_empty: root != null
 130                 return max_from(root.as(not null)).value
 131         end
 132
 133         protected fun max_from(node: N): N do
 134                 if node.right == null then return node
 135                 return max_from(node.right.as(not null))
 136         end
 137
 138         # Insert a new node in tree using `key` and `item`
 139         # O(n) in worst case, average is O(h) with h: tree height
 140         #
 141         #     var tree = new BinTreeMap[Int, String]
 142         #     tree[1] = "n1"
 143         #     assert tree.max == "n1"
 144         #     tree[3] = "n3"
 145         #     assert tree.max == "n3"
 146         redef fun []=(key, item) do
 147                 insert_node(new BinTreeNode[K, E](key, item))
 148         end
 149
 150         protected fun insert_node(node: N) do
 151                 len += 1
 152                 if root == null then
 153                         root = node
 154                 else
 155                         shift_down(root.as(not null), node)
 156                 end
 157                 if first_node == null then
 158                         first_node = node
 159                 end
 160                 if last_node != null then
 161                         last_node.next = node
 162                         node.prev = last_node
 163                 end
 164                 last_node = node
 165         end
 166
 167         # Push down the `node` in tree from a specified `from` index
 168         protected fun shift_down(from, node: N) do
 169                 if node.key < from.key then
 170                         if from.left == null then
 171                                 from.left = node
 172                                 node.parent = from
 173                         else
 174                                 shift_down(from.left.as(not null), node)
 175                         end
 176                 else if node.key > from.key then
 177                         if from.right == null then
 178                                 from.right = node
 179                                 node.parent = from
 180                         else
 181                                 shift_down(from.right.as(not null), node)
 182                         end
 183                 end
 184         end
 185
 186         # Delete node at `key` (also return the deleted node value)
 187         # O(n) in worst case, average is O(h) with h: tree height
 188         #
 189         #     var tree = new BinTreeMap[Int, String]
 190         #     tree[1] = "n1"
 191         #     assert tree.max == "n1"
 192         #     tree[3] = "n3"
 193         #     assert tree.max == "n3"
 194         #     tree.delete(3)
 195         #     assert tree.max == "n1"
 196         fun delete(key: K): nullable E do
 197                 assert is_empty: root != null
 198                 len -= 1
 199                 var node = search_down(root.as(not null), key)
 200                 if node == null then return null
 201                 if node.left == null then
 202                         transplant(node, node.right)
 203                 else if node.right == null then
 204                         transplant(node, node.left)
 205                 else
 206                         var min = min_from(node.right.as(not null))
 207                         if min.parent != node then
 208                                 transplant(min, min.right)
 209                                 min.right = node.right
 210                                 min.right.parent = min
 211                         end
 212                         transplant(node, min)
 213                         min.left = node.left
 214                         min.left.parent = min
 215                 end
 216                 if first_node == node then
 217                         first_node = null
 218                 end
 219                 if last_node == node then
 220                         last_node = node.prev
 221                         last_node.next = null
 222                 else
 223                         node.prev.next = node.next
 224                         node.next.prev = node.prev
 225                 end
 226                 return node.value
 227         end
 228
 229         # Swap a `node` with the `other` in this Tree
 230         # note: Nodes parents are updated, children still untouched
 231         protected fun transplant(node, other: nullable N) do
 232                 if node == null then return
 233                 if node.parent == null then
 234                         root = other
 235                 else if node == node.parent.left then
 236                         node.parent.left = other
 237                 else
 238                         node.parent.right = other
 239                 end
 240                 if other != null then other.parent = node.parent
 241         end
 242
 243         # Perform left rotation on `node`
 244         #
 245         #     N             Y
 246         #    / \     >     / \
 247         #   a   Y         N   c
 248         #      / \   <   / \
 249         #     b   c     a   b
 250         #
 251         protected fun rotate_left(node: N) do
 252                 var y = node.right
 253                 node.right = y.left
 254                 if y.left != null then
 255                         y.left.parent = node
 256                 end
 257                 y.parent = node.parent
 258                 if node.parent == null then
 259                         root = y
 260                 else if node == node.parent.left then
 261                         node.parent.left = y
 262                 else
 263                         node.parent.right = y
 264                 end
 265                 y.left = node
 266                 node.parent = y
 267         end
 268
 269         # Perform right rotation on `node`
 270         #
 271         #     N             Y
 272         #    / \     >     / \
 273         #   a   Y         N   c
 274         #      / \   <   / \
 275         #     b   c     a   b
 276         #
 277         protected fun rotate_right(node: N) do
 278                 var y = node.left
 279                 node.left = y.right
 280                 if y.right != null then
 281                         y.right.parent = node
 282                 end
 283                 y.parent = node.parent
 284                 if node.parent == null then
 285                         root = y
 286                 else if node == node.parent.right then
 287                         node.parent.right = y
 288                 else
 289                         node.parent.left = y
 290                 end
 291                 y.right = node
 292                 node.parent = y
 293         end
 294
 295         # Sort the tree into an array
 296         # O(n)
 297         #
 298         #     var tree = new BinTreeMap[Int, String]
 299         #     for i in [4, 2, 1, 5, 3] do tree[i] = "n{i}"
 300         #     assert tree.sort == ["n1", "n2", "n3", "n4", "n5"]
 301         fun sort: Array[E] do
 302                 var sorted = new Array[E]
 303                 if root == null then return sorted
 304                 sort_down(root.as(not null), sorted)
 305                 return sorted
 306         end
 307
 308         protected fun sort_down(node: N, sorted: Array[E]) do
 309                 if node.left != null then sort_down(node.left.as(not null), sorted)
 310                 sorted.add(node.value)
 311                 if node.right != null then sort_down(node.right.as(not null), sorted)
 312         end
 313
 314         redef fun to_s do
 315                 var root = self.root
 316                 if root == null then return "[]"
 317                 return "[{print_tree(root)}]"
 318         end
 319
 320         protected fun print_tree(node: N): String do
 321                 var s = new FlatBuffer
 322                 s.append(node.to_s)
 323                 if node.left != null then s.append(print_tree(node.left.as(not null)))
 324                 if node.right != null then s.append(print_tree(node.right.as(not null)))
 325                 return s.to_s
 326         end
 327
 328         redef fun show_dot do
 329                 assert not_empty: root != null
 330                 var f = new OProcess("dot", "-Txlib")
 331                 f.write "digraph \{\n"
 332                 dot_down(root.as(not null), f)
 333                 f.write "\}\n"
 334                 f.close
 335         end
 336
 337         protected fun dot_down(node: N, f: OProcess) do
 338                 if node.left != null then dot_down(node.left.as(not null), f)
 339                 f.write node.to_dot
 340                 if node.right != null then dot_down(node.right.as(not null), f)
 341         end
 342
 343         # O(n)
 344         #
 345         #     var tree = new BinTreeMap[Int, String]
 346         #     assert tree.length == 0
 347         #     for i in [4, 2, 1, 5, 3] do tree[i] = "n{i}"
 348         #     assert tree.length == 5
 349         redef fun length do return len
 350
 351         # Nodes are iterated in the same order in which they were added to the tree.
 352         # O(n)
 353         #
 354         #     var tree = new BinTreeMap[Int, String]
 355         #     for i in [4, 2, 1, 5, 3] do tree[i] = "n{i}"
 356         #     var keys = new Array[Int]
 357         #     for k, v in tree do
 358         #         keys.add k
 359         #     end
 360         #     assert keys == [4, 2, 1, 5, 3]
 361         redef fun iterator do return new BinTreeMapIterator[K, E](self)
 362 end
 363
 364 # TreeNode used by BinTree
 365 class BinTreeNode[K: Comparable, E]
 366         super TreeNode[K, E]
 367
 368         private var prev: nullable BinTreeNode[K, E]
 369         private var next: nullable BinTreeNode[K, E]
 370
 371         redef type SELF: BinTreeNode[K, E]
 372
 373         init(key: K, item: E) do
 374                 super(key, item)
 375         end
 376
 377         private var left_node: nullable SELF = null
 378
 379         # `left` tree node child (null if node has no left child)
 380         fun left: nullable SELF do return left_node
 381
 382         # set `left` child for this node (or null if left no child)
 383         # ENSURE: node.key < key (only if node != null)
 384         fun left=(node: nullable SELF) do
 385                 assert node != null implies node.key < key
 386                 left_node = node
 387         end
 388
 389         private var right_node: nullable SELF = null
 390
 391         # `right` tree node child (null if node has no right child)
 392         fun right: nullable SELF do return right_node
 393
 394         # set `right` child for this node (or null if right no child)
 395         # ENSURE: node.key < key (only if node != null)
 396         fun right=(node: nullable SELF) do
 397                 if node != null then
 398                         assert node.key > key
 399                 end
 400                 right_node = node
 401         end
 402
 403         # `parent` of the `parent` of this node (null if root)
 404         fun grandparent: nullable SELF do
 405                 if parent == null then
 406                         return null
 407                 else
 408                         return parent.parent
 409                 end
 410         end
 411
 412         # Other child of the `grandparent`
 413         # `left` or `right` depends on the position of the current node against its parent
 414         fun uncle: nullable SELF do
 415                 var g = grandparent
 416                 if g == null then
 417                         return null
 418                 else
 419                         if parent == g.left then
 420                                 return g.right
 421                         else
 422                                 return g.left
 423                         end
 424                 end
 425         end
 426
 427         # Other child of the parent
 428         # `left` or `right` depends on the position of the current node against its parent
 429         fun sibling: nullable SELF do
 430                 if parent == null then
 431                         return null
 432                 else if self == parent.left then
 433                         return parent.right
 434                 else if self == parent.right then
 435                         return parent.left
 436                 else
 437                         return null
 438                 end
 439         end
 440
 441         redef fun to_s do return "\{{key}: {value or else ""}\}"
 442 end
 443
 444 private class BinTreeMapIterator[K: Comparable, E]
 445         super MapIterator[K, E]
 446
 447         var current: nullable BinTreeNode[K, E]
 448
 449         init(tree: BinTreeMap[K, E]) do
 450                 current = tree.first_node
 451         end
 452
 453         redef fun is_ok do return not current == null
 454         redef fun next do current = current.next
 455         redef fun item do return current.value
 456         redef fun key do do return current.key
 457 end