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         # Get the node value associated to `key`
  48         # O(n) in worst case, average is O(h) with h: tree height
  49         #
  50         #     var tree = new BinTreeMap[Int, String]
  51         #     for i in [4, 2, 1, 5, 3] do tree[i] = "n{i}"
  52         #     assert tree[1] == "n1"
  53         redef fun [](key: K): E do
  54                 assert not_empty: root != null
  55                 var res = search_down(root.as(not null), key)
  56                 assert has_key: res != null
  57                 return res.value
  58         end
  59
  60         protected fun search_down(from: N, key: K): nullable N do
  61                 if key == from.key then return from
  62                 if from.left != null and key < from.key then
  63                         return search_down(from.left.as(not null), key)
  64                 else if from.right != null then
  65                         return search_down(from.right.as(not null), key)
  66                 end
  67                 return null
  68         end
  69
  70         # Get the node with the minimum key
  71         # O(n) in worst case, average is O(h) with h: tree height
  72         #
  73         #     var tree = new BinTreeMap[Int, String]
  74         #     for i in [4, 2, 1, 5, 3] do tree[i] = "n{i}"
  75         #     assert tree.min == "n1"
  76         fun min: E do
  77                 assert not_empty: root != null
  78                 return min_from(root.as(not null)).value
  79         end
  80
  81         protected fun min_from(node: N): N do
  82                 if node.left == null then return node
  83                 return min_from(node.left.as(not null))
  84         end
  85
  86         # Get the node with the maximum key
  87         # O(n) in worst case, average is O(h) with h: tree height
  88         #
  89         #     var tree = new BinTreeMap[Int, String]
  90         #     for i in [4, 2, 1, 5, 3, 6, 7, 8] do tree[i] = "n{i}"
  91         #     assert tree.max == "n8"
  92         fun max: E do
  93                 assert not_empty: root != null
  94                 return max_from(root.as(not null)).value
  95         end
  96
  97         protected fun max_from(node: N): N do
  98                 if node.right == null then return node
  99                 return max_from(node.right.as(not null))
 100         end
 101
 102         # Insert a new node in tree using `key` and `item`
 103         # O(n) in worst case, average is O(h) with h: tree height
 104         #
 105         #     var tree = new BinTreeMap[Int, String]
 106         #     tree[1] = "n1"
 107         #     assert tree.max == "n1"
 108         #     tree[3] = "n3"
 109         #     assert tree.max == "n3"
 110         redef fun []=(key, item) do
 111                 insert_node(new BinTreeNode[K, E](key, item))
 112         end
 113
 114         protected fun insert_node(node: N) do
 115                 if root == null then
 116                         root = node
 117                 else
 118                         shift_down(root.as(not null), node)
 119                 end
 120         end
 121
 122         # Push down the `node` in tree from a specified `from` index
 123         protected fun shift_down(from, node: N) do
 124                 if node.key < from.key then
 125                         if from.left == null then
 126                                 from.left = node
 127                                 node.parent = from
 128                         else
 129                                 shift_down(from.left.as(not null), node)
 130                         end
 131                 else if node.key > from.key then
 132                         if from.right == null then
 133                                 from.right = node
 134                                 node.parent = from
 135                         else
 136                                 shift_down(from.right.as(not null), node)
 137                         end
 138                 end
 139         end
 140
 141         # Delete node at `key` (also return the deleted node value)
 142         # O(n) in worst case, average is O(h) with h: tree height
 143         #
 144         #     var tree = new BinTreeMap[Int, String]
 145         #     tree[1] = "n1"
 146         #     assert tree.max == "n1"
 147         #     tree[3] = "n3"
 148         #     assert tree.max == "n3"
 149         #     tree.delete(3)
 150         #     assert tree.max == "n1"
 151         fun delete(key: K): nullable E do
 152                 assert is_empty: root != null
 153                 var node = search_down(root.as(not null), key)
 154                 if node == null then return null
 155                 if node.left == null then
 156                         transplant(node, node.right)
 157                 else if node.right == null then
 158                         transplant(node, node.left)
 159                 else
 160                         var min = min_from(node.right.as(not null))
 161                         if min.parent != node then
 162                                 transplant(min, min.right)
 163                                 min.right = node.right
 164                                 min.right.parent = min
 165                         end
 166                         transplant(node, min)
 167                         min.left = node.left
 168                         min.left.parent = min
 169                 end
 170                 return node.value
 171         end
 172
 173         # Swap a `node` with the `other` in this Tree
 174         # note: Nodes parents are updated, children still untouched
 175         protected fun transplant(node, other: nullable N) do
 176                 if node == null then return
 177                 if node.parent == null then
 178                         root = other
 179                 else if node == node.parent.left then
 180                         node.parent.left = other
 181                 else
 182                         node.parent.right = other
 183                 end
 184                 if other != null then other.parent = node.parent
 185         end
 186
 187         # Perform left rotation on `node`
 188         #
 189         #     N             Y
 190         #    / \     >     / \
 191         #   a   Y         N   c
 192         #      / \   <   / \
 193         #     b   c     a   b
 194         #
 195         protected fun rotate_left(node: N) do
 196                 var y = node.right
 197                 node.right = y.left
 198                 if y.left != null then
 199                         y.left.parent = node
 200                 end
 201                 y.parent = node.parent
 202                 if node.parent == null then
 203                         root = y
 204                 else if node == node.parent.left then
 205                         node.parent.left = y
 206                 else
 207                         node.parent.right = y
 208                 end
 209                 y.left = node
 210                 node.parent = y
 211         end
 212
 213         # Perform right rotation on `node`
 214         #
 215         #     N             Y
 216         #    / \     >     / \
 217         #   a   Y         N   c
 218         #      / \   <   / \
 219         #     b   c     a   b
 220         #
 221         protected fun rotate_right(node: N) do
 222                 var y = node.left
 223                 node.left = y.right
 224                 if y.right != null then
 225                         y.right.parent = node
 226                 end
 227                 y.parent = node.parent
 228                 if node.parent == null then
 229                         root = y
 230                 else if node == node.parent.right then
 231                         node.parent.right = y
 232                 else
 233                         node.parent.left = y
 234                 end
 235                 y.right = node
 236                 node.parent = y
 237         end
 238
 239         # Sort the tree into an array
 240         # O(n)
 241         #
 242         #     var tree = new BinTreeMap[Int, String]
 243         #     for i in [4, 2, 1, 5, 3] do tree[i] = "n{i}"
 244         #     assert tree.sort == ["n1", "n2", "n3", "n4", "n5"]
 245         fun sort: Array[E] do
 246                 var sorted = new Array[E]
 247                 if root == null then return sorted
 248                 sort_down(root.as(not null), sorted)
 249                 return sorted
 250         end
 251
 252         protected fun sort_down(node: N, sorted: Array[E]) do
 253                 if node.left != null then sort_down(node.left.as(not null), sorted)
 254                 sorted.add(node.value)
 255                 if node.right != null then sort_down(node.right.as(not null), sorted)
 256         end
 257
 258         redef fun to_s do
 259                 var root = self.root
 260                 if root == null then return "[]"
 261                 return "[{print_tree(root)}]"
 262         end
 263
 264         protected fun print_tree(node: N): String do
 265                 var s = new FlatBuffer
 266                 s.append(node.to_s)
 267                 if node.left != null then s.append(print_tree(node.left.as(not null)))
 268                 if node.right != null then s.append(print_tree(node.right.as(not null)))
 269                 return s.to_s
 270         end
 271
 272         redef fun show_dot do
 273                 assert not_empty: root != null
 274                 var f = new OProcess("dot", "-Txlib")
 275                 f.write "digraph \{\n"
 276                 dot_down(root.as(not null), f)
 277                 f.write "\}\n"
 278                 f.close
 279         end
 280
 281         protected fun dot_down(node: N, f: OProcess) do
 282                 if node.left != null then dot_down(node.left.as(not null), f)
 283                 f.write node.to_dot
 284                 if node.right != null then dot_down(node.right.as(not null), f)
 285         end
 286 end
 287
 288 # TreeNode used by BinTree
 289 class BinTreeNode[K: Comparable, E]
 290         super TreeNode[K, E]
 291
 292         redef type SELF: BinTreeNode[K, E]
 293
 294         init(key: K, item: E) do
 295                 super(key, item)
 296         end
 297
 298         private var left_node: nullable SELF = null
 299
 300         # `left` tree node child (null if node has no left child)
 301         fun left: nullable SELF do return left_node
 302
 303         # set `left` child for this node (or null if left no child)
 304         # ENSURE: node.key < key (only if node != null)
 305         fun left=(node: nullable SELF) do
 306                 assert node != null implies node.key < key
 307                 left_node = node
 308         end
 309
 310         private var right_node: nullable SELF = null
 311
 312         # `right` tree node child (null if node has no right child)
 313         fun right: nullable SELF do return right_node
 314
 315         # set `right` child for this node (or null if right no child)
 316         # ENSURE: node.key < key (only if node != null)
 317         fun right=(node: nullable SELF) do
 318                 if node != null then
 319                         assert node.key > key
 320                 end
 321                 right_node = node
 322         end
 323
 324         # `parent` of the `parent` of this node (null if root)
 325         fun grandparent: nullable SELF do
 326                 if parent == null then
 327                         return null
 328                 else
 329                         return parent.parent
 330                 end
 331         end
 332
 333         # Other child of the `grandparent`
 334         # `left` or `right` depends on the position of the current node against its parent
 335         fun uncle: nullable SELF do
 336                 var g = grandparent
 337                 if g == null then
 338                         return null
 339                 else
 340                         if parent == g.left then
 341                                 return g.right
 342                         else
 343                                 return g.left
 344                         end
 345                 end
 346         end
 347
 348         # Other child of the parent
 349         # `left` or `right` depends on the position of the current node against its parent
 350         fun sibling: nullable SELF do
 351                 if parent == null then
 352                         return null
 353                 else if self == parent.left then
 354                         return parent.right
 355                 else if self == parent.right then
 356                         return parent.left
 357                 else
 358                         return null
 359                 end
 360         end
 361
 362         redef fun to_s do return "\{{key}: {value or else ""}\}"
 363 end
 364