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__importabsolute_import 36 37from.treemixinimportTreeMixin 38 39__all__ = ['RBTree'] 40 41 42classNode(object): 43 """ Internal object, represents a treenode """ 44 __slots__ = ['key', 'value', 'red', 'left', 'right'] 45 46def__init__(self, key=None, value=None): 47 self.key = key 48 self.value = value 49 self.red = True 50 self.left =None51 self.right =None52 53deffree(self): 54 self.left =None55 self.right =None56 self.key =None57 self.value =None58 59def__getitem__(self, key): 60 """ x.__getitem__(key) <==> x[key], where key is 0 (left) or 1 (right) """ 61returnself.leftifkey == 0elseself.right 62 63def__setitem__(self, key, value): 64 """ x.__setitem__(key, value) <==> x[key]=value, where key is 0 (left) or 1 (right) """ 65ifkey == 0: 66 self.left = value 67else: 68 self.right = value 69 70 71defis_red(node): 72if(nodeisnotNone)andnode.red: 73returnTrue 74else: 75returnFalse 76 77 78defjsw_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 85returnsave 86 87 88defjsw_double(root, direction): 89 other_side = 1 - direction 90 root[other_side] = jsw_single(root[other_side], other_side) 91returnjsw_single(root, direction) 92 93 94classRBTree(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 """ 118def__init__(self, items=None): 119 """ x.__init__(...) initializes x; see x.__class__.__doc__ for signature """ 120 self._root =None121 self._count = 0 122ifitemsisnotNone: 123 self.update(items) 124 125defclear(self): 126 """ T.clear() -> None. Remove all items from T. """ 127def_clear(node): 128ifnodeisnotNone: 129 _clear(node.left) 130 _clear(node.right) 131 node.free() 132 _clear(self._root) 133 self._count = 0 134 self._root =None135 136 @property 137defcount(self): 138 """ count of items """ 139returnself._count 140 141 @property 142defroot(self): 143 """ root node of T """ 144returnself._root 145 146def_new_node(self, key, value): 147 """ Create a new treenode """ 148 self._count += 1 149returnNode(key, value) 150 151definsert(self, key, value): 152 """ T.insert(key, value) <==> T[key] = value, insert key, value into Tree """ 153ifself._rootisNone: # Empty tree case 154 self._root = self._new_node(key, value) 155 self._root.red = False # make root black 156return157 158 head = Node() # False tree root 159 grand_parent =None160 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 169whileTrue: 170ifnodeisNone: # Insert new node at the bottom 171 node = self._new_node(key, value) 172 parent[direction] = node 173elifis_red(node.left)andis_red(node.right): # Color flip 174 node.red = True 175 node.left.red = False 176 node.right.red = False 177 178 # Fix red violation 179ifis_red(node)andis_red(parent): 180 direction2 = 1ifgrand_grand_parent.rightisgrand_parentelse0 181ifnodeisparent[last]: 182 grand_grand_parent[direction2] = jsw_single(grand_parent, 1 - last) 183else: 184 grand_grand_parent[direction2] = jsw_double(grand_parent, 1 - last) 185 186 # Stop if found 187ifkey == node.key: 188 node.value = value # set new value for key 189break190 191 last = direction 192 direction = 0ifkey < node.keyelse1 193 # Update helpers 194ifgrand_parentisnotNone: 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 203defremove(self, key): 204 """ T.remove(key) <==> del T[key], remove item <key> from tree """ 205ifself._rootisNone: 206raiseKeyError(str(key)) 207 head = Node() # False tree root 208 node = head 209 node.right = self._root 210 parent =None211 grand_parent =None212 found =None# Found item 213 direction = 1 214 215 # Search and push a red down 216whilenode[direction]isnotNone: 217 last = direction 218 219 # Update helpers 220 grand_parent = parent 221 parent = node 222 node = node[direction] 223 224 direction = 1ifkey > node.keyelse0 225 226 # Save found node 227ifkey == node.key: 228 found = node 229 230 # Push the red node down 231ifnotis_red(node)andnotis_red(node[direction]): 232ifis_red(node[1 - direction]): 233 parent[last] = jsw_single(node, direction) 234 parent = parent[last] 235elifnotis_red(node[1 - direction]): 236 sibling = parent[1 - last] 237ifsiblingisnotNone: 238if(notis_red(sibling[1 - last]))and(notis_red(sibling[last])): 239 # Color flip 240 parent.red = False 241 sibling.red = True 242 node.red = True 243else: 244 direction2 = 1ifgrand_parent.rightisparentelse0 245ifis_red(sibling[last]): 246 grand_parent[direction2] = jsw_double(parent, last) 247elifis_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 256iffoundisnotNone: 257 found.key = node.key 258 found.value = node.value 259 parent[int(parent.rightisnode)] = node[int(node.leftisNone)] 260 node.free() 261 self._count -= 1 262 263 # Update root and make it black 264 self._root = head.right 265ifself._rootisnotNone: 266 self._root.red = False 267ifnotfound: 268raiseKeyError(str(key))