Mercurial > public > think_complexity
diff redblacktree.py @ 17:977628018b4b
The insert operation on the red-black tree seems to be working.
author | Brian Neal <bgneal@gmail.com> |
---|---|
date | Wed, 19 Dec 2012 22:29:09 -0600 |
parents | a00e97bcdb4a |
children | 92e2879e2e33 |
line wrap: on
line diff
--- a/redblacktree.py Tue Dec 18 20:03:28 2012 -0600 +++ b/redblacktree.py Wed Dec 19 22:29:09 2012 -0600 @@ -3,6 +3,12 @@ http://greenteapress.com/complexity +This code is based on the description of the red-black tree at +http://en.wikipedia.org/wiki/Red-black_tree. + +Some code and ideas were taken from code by Darren Hart at +http://dvhart.com/darren/files/rbtree.py + Copyright (C) 2012 Brian G. Neal. Permission is hereby granted, free of charge, to any person obtaining a copy of @@ -64,3 +70,210 @@ if g: return g.left if g.right is self.parent else g.right return None + + def __str__(self): + c = 'B' if self.color == BLACK else 'R' + if self.value: + return '({}: {} => {})'.format(c, self.key, self.value) + else: + return '({}: {})'.format(c, self.key) + + +class Tree(object): + """A red-black Tree class. + + A red-black tree is a binary search tree with the following properties: + + 1. A node is either red or black. + 2. The root is black. + 3. All leaves are black. + 4. Both children of every red node are black. + 5. Every simple path from a given node to any descendant leaf contains + the same number of black nodes. + + These rules ensure that the path from the root to the furthest leaf is no + more than twice as long as the path from the root to the nearest leaf. Thus + the tree is roughly height-balanced. + + """ + + def __init__(self): + self.root = None + + def __iter__(self): + return self._inorder(self.root) + + def _inorder(self, node): + """A generator to perform an inorder traversal of the nodes in the + tree starting at the given node. + + """ + if node.left: + for n in self._inorder(node.left): + yield n + + yield node + + if node.right: + for n in self._inorder(node.right): + yield n + + def insert(self, key, value=None): + node = Node(key=key, value=value, color=RED) + + # trivial case of inserting into an empty tree: + if self.root is None: + self.root = node + self.root.color = BLACK + return + + # Find a spot to insert the new red node: + + x = self.root + while x is not None: + p = x + if key < x.key: + x = x.left + else: + x = x.right + + node.parent = p # p is the new node's parent + + # p now has 1 child; decide if the new node goes on the left or right + + if key < p.key: + p.left = node + else: + p.right = node + + # ensure the new tree follows the red-black rules: + + while True: + p = node.parent + + # Case 1: root node + if p is None: + node.color = BLACK + break + + # Case 2: parent is black + if p.color == BLACK: + break + + # Case 3: parent and uncle are red + u = node.uncle + if u and u.color == RED: + # repaint parent & uncle black; grandparent becomes red: + p.color = BLACK + u.color = BLACK + gp = node.grandparent + gp.color = RED + + # enforce the rules at the grandparent level + node = gp + continue + + # gp is the grandparent; it can't be None because we know the parent + # is red at this point + gp = p.parent + + # Case 4: At this point we know the parent is red and the uncle is + # black. + # If the new node is a right child of the parent, and the + # parent is on the left of the grandparent => left rotation on the + # parent. + # If the new node is a left child of the parent, and the parent is + # on the right of the grandparent => right rotation on the parent. + + if node is p.right and p is gp.left: + self._rotate_left(p) + node = node.left + elif node is p.left and p is gp.right: + self._rotate_right(p) + node = node.right + + # Note that case 4 must fall through to case 5 to fix up the former parent + # node, which is now the child of the new red node. + + # Case 5: The parent is red and the uncle is black. + # If the new node is the left child of the parent and the parent is + # the left child of the grandparent => rotate right on the + # grandparent. + # If the new node is the right child of the parent and the parent + # is the right child of the grandparent => rotate left on the + # grandparent. + + p = node.parent + gp = p.parent + + p.color = BLACK + gp.color = RED + + if node is p.left and p is gp.left: + self._rotate_right(gp) + elif node is p.right and p is gp.right: #TODO: can this be an else? + self._rotate_left(gp) + break + + def _rotate_right(self, node): + """Rotate the tree right on the given node.""" + + p = node.parent + left = node.left + + if p: + if node is p.right: + p.right = left + else: + p.left = left + + left.parent = p + node.left = left.right + + if node.left: + node.left.parent = node + + left.right = node + node.parent = left + + # fix up the root if necessary + if self.root is node: + self.root = left + + def _rotate_left(self, node): + """Rotate the tree left on the given node.""" + + p = node.parent + right = node.right + + if p: + if node is p.right: + p.right = right + else: + p.left = right + + right.parent = p + node.right = right.left + + if node.right: + node.right.parent = node + + right.left = node + node.parent = right + + # fix up the root if necessary + if self.root is node: + self.root = right + + +if __name__ == '__main__': + import random + + tree = Tree() + + for i in range(20): + val = random.randint(0, 100) + tree.insert(val) + + for n in tree: + print n.key,