lib/trees/rbtree.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 # A red–black tree is a data structure which is a type of self-balancing binary search tree.
  16 #
  17 # Balance is preserved by painting each node of the tree with one of two colors
  18 # (typically called 'red' and 'black') in a way that satisfies certain properties,
  19 # which collectively constrain how unbalanced the tree can become in the worst case.
  20 # When the tree is modified, the new tree is subsequently rearranged and repainted
  21 # to restore the coloring properties.
  22 # The properties are designed in such a way that this rearranging and recoloring
  23 # can be performed efficiently.
  24 #
  25 # The balancing of the tree is not perfect but it is good enough to allow it
  26 # to guarantee searching in O(log n) time, where n is the total number of elements in the tree.
  27 # The insertion and deletion operations, along with the tree rearrangement and recoloring,
  28 # are also performed in O(log n) time.
  29 module rbtree
  30
  31 import bintree
  32
  33 # Red-Black Tree Map
  34 # Properties of a Red-Black tree map:
  35 # * every node is either red or black
  36 # * root is black
  37 # * every leaf (null) is black
  38 # * if a node is red, then both its children are black
  39 # * for each node, all simple path from the node to descendant
  40 #   leaves contain the same number of black nodes
  41 #
  42 # Operations:
  43 # * search average O(lg n) worst O(lg n)
  44 # * insert average O(lg n) worst O(lg n)
  45 # * delete average O(lg n) worst O(lg n)
  46 class RBTreeMap[K: Comparable, E]
  47         super BinTreeMap[K, E]
  48
  49         redef type N: RBTreeNode[K, E]
  50
  51         redef fun []=(key, item) do
  52                 insert_node(new RBTreeNode[K, E](key, item))
  53         end
  54
  55         redef fun insert_node(node) do
  56                 super
  57                 insert_fixup_case1(node)
  58         end
  59
  60         # Case 1: node is root
  61         # color node as black
  62         # it adds a black node on every path so we do nothing else
  63         private fun insert_fixup_case1(node: N) do
  64                 if node.parent == null then
  65                         node.is_red = false
  66                 else
  67                         insert_fixup_case2(node)
  68                 end
  69         end
  70
  71         # Case 2: parent is black
  72         # it do not add black node so we do nothing else
  73         private fun insert_fixup_case2(node: N) do
  74                 if node.parent.is_red then insert_fixup_case3(node)
  75         end
  76
  77         # Case 3: node, parent and uncle are red
  78         # this is a LLr or RRr conflict
  79         # we recolor recursively the tree to the root
  80         private fun insert_fixup_case3(node: N) do
  81                 if node.uncle != null and node.uncle.is_red then
  82                         node.parent.is_red = false
  83                         node.uncle.is_red = false
  84                         node.grandparent.is_red = true
  85                         insert_fixup_case1(node.grandparent.as(not null))
  86                 else
  87                         insert_fixup_case4(node)
  88                 end
  89         end
  90
  91         # Case 4: node and parent are red, uncle is black
  92         # this is a LRb or RLb conflict
  93         # we rotate the tree to balance it
  94         private fun insert_fixup_case4(node: N) do
  95                 if node == node.parent.right and node.parent == node.grandparent.left then
  96                         rotate_left(node.parent.as(not null))
  97                         node = node.left.as(not null)
  98                 else if node == node.parent.left and node.parent == node.grandparent.right then
  99                         rotate_right(node.parent.as(not null))
 100                         node = node.right.as(not null)
 101                 end
 102                 insert_fixup_case5(node)
 103         end
 104
 105         # Case 5: node and parent are red, uncle is black
 106         # this is a LLb or RRb conflict
 107         # we rotate the tree to balance it
 108         private fun insert_fixup_case5(node: N) do
 109                 node.parent.is_red = false
 110                 node.grandparent.is_red = true
 111                 if node == node.parent.left then
 112                         rotate_right(node.grandparent.as(not null))
 113                 else
 114                         rotate_left(node.grandparent.as(not null))
 115                 end
 116         end
 117
 118         # TODO implement RBTree::delete
 119         redef fun delete(key) is abstract
 120
 121         redef fun dot_down(node, f) do
 122                 if node.left != null then dot_down(node.left.as(not null), f)
 123                 f.write node.to_dot
 124                 if node.parent != null then f.write "\"{node.parent.as(not null)}\" -> \"{node}\"[dir=back];\n"
 125                 if node.right != null then dot_down(node.right.as(not null), f)
 126         end
 127 end
 128
 129 # RedBlackTree node (can be red or black)
 130 class RBTreeNode[K: Comparable, E]
 131         super BinTreeNode[K, E]
 132
 133         redef type N: RBTreeNode[K, E]
 134
 135         # Is the node red?
 136         private var is_red = true
 137
 138         redef fun to_dot do
 139                 if is_red then
 140                         return "\"{self}\"[style=filled,fillcolor=red,fontcolor=white];\n"
 141                 else
 142                         return "\"{self}\"[style=filled,fillcolor=black,fontcolor=white];\n"
 143                 end
 144
 145         end
 146 end
 147