src/third_party/bintrees/bintrees/rbtree.py

   1 #!/usr/bin/env python
   2 #coding:utf-8
   3 # Author:  mozman (python version)
   4 # Purpose: red-black tree module (Julienne Walker's none recursive algorithm)
   5 # source: http://eternallyconfuzzled.com/tuts/datastructures/jsw_tut_rbtree.aspx
   6 # Created: 01.05.2010
   7 # Copyright (c) 2010-2013 by Manfred Moitzi
   8 # License: MIT License
   9
  10 # Conclusion of Julian Walker
  11
  12 # Red black trees are interesting beasts. They're believed to be simpler than
  13 # AVL trees (their direct competitor), and at first glance this seems to be the
  14 # case because insertion is a breeze. However, when one begins to play with the
  15 # deletion algorithm, red black trees become very tricky. However, the
  16 # counterweight to this added complexity is that both insertion and deletion
  17 # can be implemented using a single pass, top-down algorithm. Such is not the
  18 # case with AVL trees, where only the insertion algorithm can be written top-down.
  19 # Deletion from an AVL tree requires a bottom-up algorithm.
  20
  21 # So when do you use a red black tree? That's really your decision, but I've
  22 # found that red black trees are best suited to largely random data that has
  23 # occasional degenerate runs, and searches have no locality of reference. This
  24 # takes full advantage of the minimal work that red black trees perform to
  25 # maintain balance compared to AVL trees and still allows for speedy searches.
  26
  27 # Red black trees are popular, as most data structures with a whimsical name.
  28 # For example, in Java and C++, the library map structures are typically
  29 # implemented with a red black tree. Red black trees are also comparable in
  30 # speed to AVL trees. While the balance is not quite as good, the work it takes
  31 # to maintain balance is usually better in a red black tree. There are a few
  32 # misconceptions floating around, but for the most part the hype about red black
  33 # trees is accurate.
  34
  35 from __future__ import absolute_import
  36
  37 from .treemixin import TreeMixin
  38
  39 __all__ = ['RBTree']
  40
  41
  42 class Node(object):
  43     """ Internal object, represents a treenode """
  44     __slots__ = ['key', 'value', 'red', 'left', 'right']
  45
  46     def __init__(self, key=None, value=None):
  47         self.key = key
  48         self.value = value
  49         self.red = True
  50         self.left = None
  51         self.right = None
  52
  53     def free(self):
  54         self.left = None
  55         self.right = None
  56         self.key = None
  57         self.value = None
  58
  59     def __getitem__(self, key):
  60         """ x.__getitem__(key) <==> x[key], where key is 0 (left) or 1 (right) """
  61         return self.left if key == 0 else self.right
  62
  63     def __setitem__(self, key, value):
  64         """ x.__setitem__(key, value) <==> x[key]=value, where key is 0 (left) or 1 (right) """
  65         if key == 0:
  66             self.left = value
  67         else:
  68             self.right = value
  69
  70
  71 def is_red(node):
  72     if (node is not None) and node.red:
  73         return True
  74     else:
  75         return False
  76
  77
  78 def jsw_single(root, direction):
  79     other_side = 1 - direction
  80     save = root[other_side]
  81     root[other_side] = save[direction]
  82     save[direction] = root
  83     root.red = True
  84     save.red = False
  85     return save
  86
  87
  88 def jsw_double(root, direction):
  89     other_side = 1 - direction
  90     root[other_side] = jsw_single(root[other_side], other_side)
  91     return jsw_single(root, direction)
  92
  93
  94 class RBTree(TreeMixin):
  95     """
  96     RBTree implements a balanced binary tree with a dict-like interface.
  97
  98     see: http://en.wikipedia.org/wiki/Red_black_tree
  99
 100     A red-black tree is a type of self-balancing binary search tree, a data
 101     structure used in computing science, typically used to implement associative
 102     arrays. The original structure was invented in 1972 by Rudolf Bayer, who
 103     called them "symmetric binary B-trees", but acquired its modern name in a
 104     paper in 1978 by Leonidas J. Guibas and Robert Sedgewick. It is complex,
 105     but has good worst-case running time for its operations and is efficient in
 106     practice: it can search, insert, and delete in O(log n) time, where n is
 107     total number of elements in the tree. Put very simply, a red-black tree is a
 108     binary search tree which inserts and removes intelligently, to ensure the
 109     tree is reasonably balanced.
 110
 111     RBTree() -> new empty tree.
 112     RBTree(mapping) -> new tree initialized from a mapping
 113     RBTree(seq) -> new tree initialized from seq [(k1, v1), (k2, v2), ... (kn, vn)]
 114
 115     see also TreeMixin() class.
 116
 117     """
 118     def __init__(self, items=None):
 119         """ x.__init__(...) initializes x; see x.__class__.__doc__ for signature """
 120         self._root = None
 121         self._count = 0
 122         if items is not None:
 123             self.update(items)
 124
 125     def clear(self):
 126         """ T.clear() -> None.  Remove all items from T. """
 127         def _clear(node):
 128             if node is not None:
 129                 _clear(node.left)
 130                 _clear(node.right)
 131                 node.free()
 132         _clear(self._root)
 133         self._count = 0
 134         self._root = None
 135
 136     @property
 137     def count(self):
 138         """ count of items """
 139         return self._count
 140
 141     @property
 142     def root(self):
 143         """ root node of T """
 144         return self._root
 145
 146     def _new_node(self, key, value):
 147         """ Create a new treenode """
 148         self._count += 1
 149         return Node(key, value)
 150
 151     def insert(self, key, value):
 152         """ T.insert(key, value) <==> T[key] = value, insert key, value into Tree """
 153         if self._root is None:  # Empty tree case
 154             self._root = self._new_node(key, value)
 155             self._root.red = False  # make root black
 156             return
 157
 158         head = Node()  # False tree root
 159         grand_parent = None
 160         grand_grand_parent = head
 161         parent = None  # parent
 162         direction = 0
 163         last = 0
 164
 165         # Set up helpers
 166         grand_grand_parent.right = self._root
 167         node = grand_grand_parent.right
 168         # Search down the tree
 169         while True:
 170             if node is None:  # Insert new node at the bottom
 171                 node = self._new_node(key, value)
 172                 parent[direction] = node
 173             elif is_red(node.left) and is_red(node.right):  # Color flip
 174                 node.red = True
 175                 node.left.red = False
 176                 node.right.red = False
 177
 178             # Fix red violation
 179             if is_red(node) and is_red(parent):
 180                 direction2 = 1 if grand_grand_parent.right is grand_parent else 0
 181                 if node is parent[last]:
 182                     grand_grand_parent[direction2] = jsw_single(grand_parent, 1 - last)
 183                 else:
 184                     grand_grand_parent[direction2] = jsw_double(grand_parent, 1 - last)
 185
 186             # Stop if found
 187             if key == node.key:
 188                 node.value = value  # set new value for key
 189                 break
 190
 191             last = direction
 192             direction = 0 if key < node.key else 1
 193             # Update helpers
 194             if grand_parent is not None:
 195                 grand_grand_parent = grand_parent
 196             grand_parent = parent
 197             parent = node
 198             node = node[direction]
 199
 200         self._root = head.right  # Update root
 201         self._root.red = False  # make root black
 202
 203     def remove(self, key):
 204         """ T.remove(key) <==> del T[key], remove item <key> from tree """
 205         if self._root is None:
 206             raise KeyError(str(key))
 207         head = Node()  # False tree root
 208         node = head
 209         node.right = self._root
 210         parent = None
 211         grand_parent = None
 212         found = None  # Found item
 213         direction = 1
 214
 215         # Search and push a red down
 216         while node[direction] is not None:
 217             last = direction
 218
 219             # Update helpers
 220             grand_parent = parent
 221             parent = node
 222             node = node[direction]
 223
 224             direction = 1 if key > node.key else 0
 225
 226             # Save found node
 227             if key == node.key:
 228                 found = node
 229
 230             # Push the red node down
 231             if not is_red(node) and not is_red(node[direction]):
 232                 if is_red(node[1 - direction]):
 233                     parent[last] = jsw_single(node, direction)
 234                     parent = parent[last]
 235                 elif not is_red(node[1 - direction]):
 236                     sibling = parent[1 - last]
 237                     if sibling is not None:
 238                         if (not is_red(sibling[1 - last])) and (not is_red(sibling[last])):
 239                             # Color flip
 240                             parent.red = False
 241                             sibling.red = True
 242                             node.red = True
 243                         else:
 244                             direction2 = 1 if grand_parent.right is parent else 0
 245                             if is_red(sibling[last]):
 246                                 grand_parent[direction2] = jsw_double(parent, last)
 247                             elif is_red(sibling[1-last]):
 248                                 grand_parent[direction2] = jsw_single(parent, last)
 249                             # Ensure correct coloring
 250                             grand_parent[direction2].red = True
 251                             node.red = True
 252                             grand_parent[direction2].left.red = False
 253                             grand_parent[direction2].right.red = False
 254
 255         # Replace and remove if found
 256         if found is not None:
 257             found.key = node.key
 258             found.value = node.value
 259             parent[int(parent.right is node)] = node[int(node.left is None)]
 260             node.free()
 261             self._count -= 1
 262
 263         # Update root and make it black
 264         self._root = head.right
 265         if self._root is not None:
 266             self._root.red = False
 267         if not found:
 268             raise KeyError(str(key))