# HG changeset patch
# User Brian Neal <bgneal@gmail.com>
# Date 1355977749 21600
# Node ID 977628018b4bdfbb63ea2a279218d82df9120b1f
# Parent  a00e97bcdb4a80dc8b9a4ee75e88317d191b0889
The insert operation on the red-black tree seems to be working.

diff -r a00e97bcdb4a -r 977628018b4b redblacktree.py
--- 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,