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