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 35from __future__ import absolute_import 36 37from .treemixin import TreeMixin 38 39__all__ = ['RBTree'] 40 41 42class 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 71def is_red(node): 72 if (node is not None) and node.red: 73 return True 74 else: 75 return False 76 77 78def 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 88def 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 94class 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))