"""Red Black Tree A binary search tree becomes less efficient when the tree becomes unbalanced. For instance, if we store a set of numbers in increasing order, we get a tree that essentially is a linked list with a lot of empty left-child pointers. """ import binary_search_tree as bintree ###################################################################### # Constants UNDEFINED = 0 RED = 10 BLACK = 20 ###################################################################### # Data Structures class RedBlackTreeNode: def __init__(self, key, nil, color=UNDEFINED): self.key = key self.left = nil self.right = nil self.parent = nil self.color = color class RedBlackTree: def __init__(self, root=None): self.__init_nil() if root == None: self.root = self.nil else: self.root = root def __init_nil(self): self.nil = RedBlackTreeNode(None, None) self.nil.left = self.nil self.nil.right = self.nil self.nil.parent = self.nil self.nil.color = BLACK ###################################################################### # Public def insert(T, x): bintree.insert(T, x, T.nil) x.color = RED while x.parent.color == RED: if x.parent == x.parent.parent.left: y = x.parent.parent.right if y.color == RED: x.parent.color = BLACK y.color = BLACK x.parent.parent.color = RED x = x.parent.parent else: if x == x.parent.right: x = x.parent __left_rotate(T, x) x.parent.color = BLACK x.parent.parent.color = RED __right_rotate(T, x.parent.parent) else: y = x.parent.parent.left if y.color == RED: x.parent.color = BLACK y.color = BLACK x.parent.parent.color = RED x = x.parent.parent else: if x == x.parent.left: x = x.parent __right_rotate(T, x) x.parent.color = BLACK x.parent.parent.color = RED __left_rotate(T, x.parent.parent) T.root.color = BLACK def delete(T, z): if z.left == T.nil or z.right == T.nil: y = z else: y = bintree.successor(z, T.nil) if y.left != T.nil: x = y.left else: x = y.right x.parent = y.parent if y.parent == T.nil: T.root = x else: if y == y.parent.left: y.parent.left = x else: y.parent.right = x if y != z: z.key = y.key if y.color == BLACK: __delete_fixup(T, x) return y def search(T, k): return bintree.iterative_search(T.root, k, T.nil) ###################################################################### # Private def __delete_fixup(T, x): while x != T.root and x.color == BLACK: if x == x.parent.left: w = x.parent.right if w.color == RED: w.color = BLACK x.parent.color = RED __left_rotate(T, x.parent) w = x.parent.right if w.left.color == BLACK and w.right.color == BLACK: w.color = RED x = x.parent else: if w.right.color == BLACK: w.left.color = BLACK w.color = RED __right_rotate(T, w) w = x.parent.right w.color = x.parent.color x.parent.color = BLACK w.right.color = BLACK __left_rotate(T, x.parent) x = T.root else: w = x.parent.left if w.color == RED: w.color = BLACK x.parent.color = RED __right_rotate(T, x.parent) w = x.parent.left if w.right.color == BLACK and w.left.color == BLACK: w.color = RED x = x.parent else: if w.left.color == BLACK: w.right.color = BLACK w.color = RED __left_rotate(T, w) w = x.parent.left w.color = x.parent.color x.parent.color = BLACK w.left.color = BLACK __right_rotate(T, x.parent) x = T.root x.color = BLACK def __left_rotate(T, x): y = x.right x.right = y.left y.left.parent = x y.parent = x.parent if x.parent == T.nil: T.root = y else: if x == x.parent.left: x.parent.left = y else: x.parent.right = y y.left = x x.parent = y def __right_rotate(T, x): y = x.left x.left = y.right y.right.parent = x y.parent = x.parent if x.parent == T.nil: T.root = y else: if x == x.parent.right: x.parent.right = y else: x.parent.left = y y.right = x x.parent = y ###################################################################### # TEST def test(): data = [12,6,14,17,15,16,10,2,18,19] T = RedBlackTree() for d in data: z = RedBlackTreeNode(d, T.nil) insert(T, z) print "Insert %2d:"%d, bintree.print_tree(T, T.nil) for i in xrange(0, len(data), 2): test_search_and_delete(T, data[i]) def test_search(T, k): print "Searching for %d..."%(k), z = search(T, k) print z print "Node key=%d"%(z.key) return z def test_delete(T, z): print "Deleting", z delete(T, z) bintree.print_tree(T, T.nil) def test_search_and_delete(T, k): z = test_search(T, k) test_delete(T, z) if __name__ == '__main__': test()