--- a/redblacktree.py	Wed Dec 19 22:29:09 2012 -0600
+++ b/redblacktree.py	Wed Dec 26 19:59:17 2012 -0600
@@ -1,14 +1,4 @@
-"""A red-black tree for Section 3.4, Exercise 4 in Allen Downey's _Think
-Complexity_ book.
-
-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
@@ -28,48 +18,58 @@
 IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 
+----
+
+A red-black tree for Section 3.4, Exercise 4 in Allen Downey's _Think
+Complexity_ book.
+
+http://greenteapress.com/complexity
+
+Red black trees are described on Wikipedia:
+http://en.wikipedia.org/wiki/Red-black_tree.
+
+This is basically a Python implementation of Julienne Walker's Red Black Trees
+tutorial found at:
+http://www.eternallyconfuzzled.com/tuts/datastructures/jsw_tut_rbtree.aspx
+We implement Julienne's top-down insertion and deletion algorithms here.
+
+Some ideas were also borrowed from code by Darren Hart at
+http://dvhart.com/darren/files/rbtree.py
+
 """
 
 BLACK, RED = range(2)
+LEFT, RIGHT = range(2)
+
+
+class TreeError(Exception):
+    """Base exception class for red-black tree errors."""
+    pass
+
 
 class Node(object):
     """A node class for red-black trees.
 
-    A node has an optional parent, and optional left and right children. Each
-    node also has a color, either red or black. A node has a key and an optional
-    value. The key is used to order the red black tree by calling the "<"
-    operator when comparing keys. The optional value is useful for using the
-    red-black tree to implement a map datastructure.
+    * A node has a color, either RED or BLACK.
+    * A node has a key and an optional value.
 
-    In a red-black tree, nil children are always considered black.
+        * The key is used to order the red-black tree by calling the "<"
+          operator when comparing two nodes.
+        * The value is useful for using the red-black tree to implement a map
+          datastructure.
+
+    * Nodes have exactly 2 link fields which we represent as a list of
+      2 elements to represent the left and right children of the node. This list
+      representation was borrowed from Julienne Walker as it simplifies some of
+      the code. Element 0 is the LEFT child and element 1 is the RIGHT child.
 
     """
 
-    def __init__(self, key, color=RED, value=None):
+    def __init__(self, key, value=None, color=RED):
         self.key = key
         self.value = value
         self.color = color
-
-        self.parent = None
-        self.left = None
-        self.right = None
-
-    @property
-    def grandparent(self):
-        """Return the grandparent of this node, or None if it does not exist."""
-        return self.parent.parent if self.parent else None
-
-    @property
-    def uncle(self):
-        """Return this node's uncle if it exists, or None if not.
-
-        An uncle is a parent's sibling.
-
-        """
-        g = self.grandparent
-        if g:
-            return g.left if g.right is self.parent else g.right
-        return None
+        self.link = [None, None]
 
     def __str__(self):
         c = 'B' if self.color == BLACK else 'R'
@@ -79,6 +79,16 @@
             return '({}: {})'.format(c, self.key)
 
 
+def is_red(n):
+    """Return True if the given Node n is RED and False otherwise.
+
+    If the node is None, then it is considered to be BLACK, and False is
+    returned.
+
+    """
+    return n is not None and n.color == RED
+
+
 class Tree(object):
     """A red-black Tree class.
 
@@ -108,172 +118,161 @@
         tree starting at the given node.
 
         """
-        if node.left:
-            for n in self._inorder(node.left):
+        if node.link[0]:
+            for n in self._inorder(node.link[0]):
                 yield n
 
         yield node
 
-        if node.right:
-            for n in self._inorder(node.right):
+        if node.link[1]:
+            for n in self._inorder(node.link[1]):
                 yield n
 
+    def _single_rotate(self, root, d):
+        """Perform a single rotation about the node 'root' in the given
+        direction 'd' (LEFT or RIGHT).
+
+        The old root is set to RED and the new root is set to BLACK.
+
+        Returns the new root.
+
+        """
+        nd = not d
+        save = root.link[nd]
+
+        root.link[nd] = save.link[d]
+        save.link[d] = root
+
+        root.color = RED
+        save.color = BLACK
+
+        return save
+
+    def _double_rotate(self, root, d):
+        """Perform two single rotations about the node root in the direction d.
+
+        The new root is returned.
+
+        """
+        nd = not d
+        root.link[nd] = self._single_rotate(root.link[nd], nd)
+        return self._single_rotate(root, d)
+
+    def validate(self, root=None):
+        """Checks to see if the red-black tree validates at the given node, i.e.
+        all red-black tree rules are valid. If root is None, the root of the
+        tree is used as the starting point.
+
+        If any rules are violated, a TreeError is raised.
+
+        Returns the black height of the tree rooted at root.
+
+        """
+        if root is None:
+            root = self.root
+
+        return self._validate(root)
+
+    def _validate(self, root):
+        """Internal implementation of the validate() method."""
+
+        if root is None:
+            return 1
+
+        ln = root.link[0]
+        rn = root.link[1]
+
+        # Check for consecutive red links
+
+        if is_red(root) and (is_red(ln) or is_red(rn)):
+            raise TreeError('red violation')
+
+        lh = self._validate(ln)
+        rh = self._validate(rn)
+
+        # Check for invalid binary search tree
+
+        if (ln and ln.key >= root.key) or (rn and rn.key <= root.key):
+            raise TreeError('binary tree violation')
+
+        # Check for black height mismatch
+        if lh != rh:
+            raise TreeError('black violation')
+
+        # Only count black links
+
+        return lh if is_red(root) else lh + 1
+
     def insert(self, key, value=None):
-        node = Node(key=key, value=value, color=RED)
+        """Insert the (key, value) pair into the tree."""
 
-        # trivial case of inserting into an empty tree:
+        # Check for the empty tree case:
         if self.root is None:
-            self.root = node
-            self.root.color = BLACK
+            self.root = Node(key=key, value=value, color=BLACK)
             return
 
-        # Find a spot to insert the new red node:
+        # False/dummy tree root
+        head = Node(key=None, value=None, color=BLACK)
+        d = LEFT
+        last = LEFT
 
-        x = self.root
-        while x is not None:
-            p = x
-            if key < x.key:
-                x = x.left
-            else:
-                x = x.right
+        t = head
+        g = p = None
+        t.link[1] = self.root
+        q = self.root
 
-        node.parent = p     # p is the new node's parent
+        # Search down the tree
+        while True:
+            if q is None:
+                # Insert new node at the bottom
+                p.link[d] = q = Node(key=key, value=value, color=RED)
+            elif is_red(q.link[0]) and is_red(q.link[1]):
+                # Color flip
+                q.color = RED
+                q.link[0].color = BLACK
+                q.link[1].color = BLACK
 
-        # p now has 1 child; decide if the new node goes on the left or right
+            # Fix red violation
+            if is_red(q) and is_red(p):
+                d2 = t.link[1] is g
 
-        if key < p.key:
-            p.left = node
-        else:
-            p.right = node
+                if q is p.link[last]:
+                    t.link[d2] = self._single_rotate(g, not last)
+                else:
+                    t.link[d2] = self._double_rotate(g, not last)
 
-        # 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
+            # Stop if found
+            if q.key == key:
                 break
 
-            # Case 2: parent is black
-            if p.color == BLACK:
-                break
+            last = d
+            d = q.key < key
 
-            # 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
+            # Update helpers
+            if g is not None:
+                t = g;
+            g = p
+            p = q
+            q = q.link[d]
 
-                # 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
+        # Update root
+        self.root = head.link[1]
+        self.root.color = BLACK
 
 
 if __name__ == '__main__':
     import random
 
-    tree = Tree()
+    for n in range(30):
+        tree = Tree()
+        tree.validate()
 
-    for i in range(20):
-        val = random.randint(0, 100)
-        tree.insert(val)
+        for i in range(25):
+            val = random.randint(0, 100)
+            tree.insert(val)
 
-    for n in tree:
-        print n.key,
+        for n in tree:
+            print n.key,
+        print
+
+        tree.validate()