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