bgneal@18: """
bgneal@15: Copyright (C) 2012 Brian G. Neal.
bgneal@15: 
bgneal@15: Permission is hereby granted, free of charge, to any person obtaining a copy of
bgneal@15: this software and associated documentation files (the "Software"), to deal in
bgneal@15: the Software without restriction, including without limitation the rights to
bgneal@15: use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
bgneal@15: the Software, and to permit persons to whom the Software is furnished to do so,
bgneal@15: subject to the following conditions:
bgneal@15: 
bgneal@15: The above copyright notice and this permission notice shall be included in all
bgneal@15: copies or substantial portions of the Software.
bgneal@15: 
bgneal@15: THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
bgneal@15: IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
bgneal@15: FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
bgneal@15: COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
bgneal@15: IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
bgneal@15: CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
bgneal@15: 
bgneal@18: ----
bgneal@18: 
bgneal@18: A red-black tree for Section 3.4, Exercise 4 in Allen Downey's _Think
bgneal@18: Complexity_ book.
bgneal@18: 
bgneal@18: http://greenteapress.com/complexity
bgneal@18: 
bgneal@18: Red black trees are described on Wikipedia:
bgneal@18: http://en.wikipedia.org/wiki/Red-black_tree.
bgneal@18: 
bgneal@18: This is basically a Python implementation of Julienne Walker's Red Black Trees
bgneal@18: tutorial found at:
bgneal@18: http://www.eternallyconfuzzled.com/tuts/datastructures/jsw_tut_rbtree.aspx
bgneal@18: We implement Julienne's top-down insertion and deletion algorithms here.
bgneal@18: 
bgneal@18: Some ideas were also borrowed from code by Darren Hart at
bgneal@18: http://dvhart.com/darren/files/rbtree.py
bgneal@18: 
bgneal@15: """
bgneal@15: 
bgneal@16: BLACK, RED = range(2)
bgneal@18: LEFT, RIGHT = range(2)
bgneal@18: 
bgneal@18: 
bgneal@18: class TreeError(Exception):
bgneal@18:     """Base exception class for red-black tree errors."""
bgneal@18:     pass
bgneal@18: 
bgneal@15: 
bgneal@15: class Node(object):
bgneal@15:     """A node class for red-black trees.
bgneal@15: 
bgneal@18:     * A node has a color, either RED or BLACK.
bgneal@18:     * A node has a key and an optional value.
bgneal@15: 
bgneal@18:         * The key is used to order the red-black tree by calling the "<"
bgneal@18:           operator when comparing two nodes.
bgneal@18:         * The value is useful for using the red-black tree to implement a map
bgneal@18:           datastructure.
bgneal@18: 
bgneal@18:     * Nodes have exactly 2 link fields which we represent as a list of
bgneal@18:       2 elements to represent the left and right children of the node. This list
bgneal@18:       representation was borrowed from Julienne Walker as it simplifies some of
bgneal@18:       the code. Element 0 is the LEFT child and element 1 is the RIGHT child.
bgneal@15: 
bgneal@15:     """
bgneal@15: 
bgneal@18:     def __init__(self, key, value=None, color=RED):
bgneal@15:         self.key = key
bgneal@15:         self.value = value
bgneal@15:         self.color = color
bgneal@18:         self.link = [None, None]
bgneal@17: 
bgneal@17:     def __str__(self):
bgneal@17:         c = 'B' if self.color == BLACK else 'R'
bgneal@17:         if self.value:
bgneal@17:             return '({}: {} => {})'.format(c, self.key, self.value)
bgneal@17:         else:
bgneal@17:             return '({}: {})'.format(c, self.key)
bgneal@17: 
bgneal@17: 
bgneal@18: def is_red(n):
bgneal@18:     """Return True if the given Node n is RED and False otherwise.
bgneal@18: 
bgneal@18:     If the node is None, then it is considered to be BLACK, and False is
bgneal@18:     returned.
bgneal@18: 
bgneal@18:     """
bgneal@18:     return n is not None and n.color == RED
bgneal@18: 
bgneal@18: 
bgneal@17: class Tree(object):
bgneal@17:     """A red-black Tree class.
bgneal@17: 
bgneal@17:     A red-black tree is a binary search tree with the following properties:
bgneal@17: 
bgneal@17:         1. A node is either red or black.
bgneal@17:         2. The root is black.
bgneal@17:         3. All leaves are black.
bgneal@17:         4. Both children of every red node are black.
bgneal@17:         5. Every simple path from a given node to any descendant leaf contains
bgneal@17:            the same number of black nodes.
bgneal@17: 
bgneal@17:     These rules ensure that the path from the root to the furthest leaf is no
bgneal@17:     more than twice as long as the path from the root to the nearest leaf. Thus
bgneal@17:     the tree is roughly height-balanced.
bgneal@17: 
bgneal@17:     """
bgneal@17: 
bgneal@17:     def __init__(self):
bgneal@17:         self.root = None
bgneal@17: 
bgneal@17:     def __iter__(self):
bgneal@17:         return self._inorder(self.root)
bgneal@17: 
bgneal@17:     def _inorder(self, node):
bgneal@17:         """A generator to perform an inorder traversal of the nodes in the
bgneal@17:         tree starting at the given node.
bgneal@17: 
bgneal@17:         """
bgneal@18:         if node.link[0]:
bgneal@18:             for n in self._inorder(node.link[0]):
bgneal@17:                 yield n
bgneal@17: 
bgneal@17:         yield node
bgneal@17: 
bgneal@18:         if node.link[1]:
bgneal@18:             for n in self._inorder(node.link[1]):
bgneal@17:                 yield n
bgneal@17: 
bgneal@18:     def _single_rotate(self, root, d):
bgneal@18:         """Perform a single rotation about the node 'root' in the given
bgneal@18:         direction 'd' (LEFT or RIGHT).
bgneal@18: 
bgneal@18:         The old root is set to RED and the new root is set to BLACK.
bgneal@18: 
bgneal@18:         Returns the new root.
bgneal@18: 
bgneal@18:         """
bgneal@18:         nd = not d
bgneal@18:         save = root.link[nd]
bgneal@18: 
bgneal@18:         root.link[nd] = save.link[d]
bgneal@18:         save.link[d] = root
bgneal@18: 
bgneal@18:         root.color = RED
bgneal@18:         save.color = BLACK
bgneal@18: 
bgneal@18:         return save
bgneal@18: 
bgneal@18:     def _double_rotate(self, root, d):
bgneal@18:         """Perform two single rotations about the node root in the direction d.
bgneal@18: 
bgneal@18:         The new root is returned.
bgneal@18: 
bgneal@18:         """
bgneal@18:         nd = not d
bgneal@18:         root.link[nd] = self._single_rotate(root.link[nd], nd)
bgneal@18:         return self._single_rotate(root, d)
bgneal@18: 
bgneal@18:     def validate(self, root=None):
bgneal@18:         """Checks to see if the red-black tree validates at the given node, i.e.
bgneal@18:         all red-black tree rules are valid. If root is None, the root of the
bgneal@18:         tree is used as the starting point.
bgneal@18: 
bgneal@18:         If any rules are violated, a TreeError is raised.
bgneal@18: 
bgneal@18:         Returns the black height of the tree rooted at root.
bgneal@18: 
bgneal@18:         """
bgneal@18:         if root is None:
bgneal@18:             root = self.root
bgneal@18: 
bgneal@18:         return self._validate(root)
bgneal@18: 
bgneal@18:     def _validate(self, root):
bgneal@18:         """Internal implementation of the validate() method."""
bgneal@18: 
bgneal@18:         if root is None:
bgneal@18:             return 1
bgneal@18: 
bgneal@18:         ln = root.link[0]
bgneal@18:         rn = root.link[1]
bgneal@18: 
bgneal@18:         # Check for consecutive red links
bgneal@18: 
bgneal@18:         if is_red(root) and (is_red(ln) or is_red(rn)):
bgneal@18:             raise TreeError('red violation')
bgneal@18: 
bgneal@18:         lh = self._validate(ln)
bgneal@18:         rh = self._validate(rn)
bgneal@18: 
bgneal@18:         # Check for invalid binary search tree
bgneal@18: 
bgneal@18:         if (ln and ln.key >= root.key) or (rn and rn.key <= root.key):
bgneal@18:             raise TreeError('binary tree violation')
bgneal@18: 
bgneal@18:         # Check for black height mismatch
bgneal@18:         if lh != rh:
bgneal@18:             raise TreeError('black violation')
bgneal@18: 
bgneal@18:         # Only count black links
bgneal@18: 
bgneal@18:         return lh if is_red(root) else lh + 1
bgneal@18: 
bgneal@17:     def insert(self, key, value=None):
bgneal@18:         """Insert the (key, value) pair into the tree."""
bgneal@17: 
bgneal@18:         # Check for the empty tree case:
bgneal@17:         if self.root is None:
bgneal@18:             self.root = Node(key=key, value=value, color=BLACK)
bgneal@17:             return
bgneal@17: 
bgneal@18:         # False/dummy tree root
bgneal@18:         head = Node(key=None, value=None, color=BLACK)
bgneal@18:         d = LEFT
bgneal@18:         last = LEFT
bgneal@17: 
bgneal@18:         t = head
bgneal@18:         g = p = None
bgneal@18:         t.link[1] = self.root
bgneal@18:         q = self.root
bgneal@17: 
bgneal@18:         # Search down the tree
bgneal@18:         while True:
bgneal@18:             if q is None:
bgneal@18:                 # Insert new node at the bottom
bgneal@18:                 p.link[d] = q = Node(key=key, value=value, color=RED)
bgneal@18:             elif is_red(q.link[0]) and is_red(q.link[1]):
bgneal@18:                 # Color flip
bgneal@18:                 q.color = RED
bgneal@18:                 q.link[0].color = BLACK
bgneal@18:                 q.link[1].color = BLACK
bgneal@17: 
bgneal@18:             # Fix red violation
bgneal@18:             if is_red(q) and is_red(p):
bgneal@18:                 d2 = t.link[1] is g
bgneal@17: 
bgneal@18:                 if q is p.link[last]:
bgneal@18:                     t.link[d2] = self._single_rotate(g, not last)
bgneal@18:                 else:
bgneal@18:                     t.link[d2] = self._double_rotate(g, not last)
bgneal@17: 
bgneal@18:             # Stop if found
bgneal@18:             if q.key == key:
bgneal@17:                 break
bgneal@17: 
bgneal@18:             last = d
bgneal@18:             d = q.key < key
bgneal@17: 
bgneal@18:             # Update helpers
bgneal@18:             if g is not None:
bgneal@18:                 t = g;
bgneal@18:             g = p
bgneal@18:             p = q
bgneal@18:             q = q.link[d]
bgneal@17: 
bgneal@18:         # Update root
bgneal@18:         self.root = head.link[1]
bgneal@18:         self.root.color = BLACK
bgneal@17: 
bgneal@17: 
bgneal@17: if __name__ == '__main__':
bgneal@17:     import random
bgneal@17: 
bgneal@18:     for n in range(30):
bgneal@18:         tree = Tree()
bgneal@18:         tree.validate()
bgneal@17: 
bgneal@18:         for i in range(25):
bgneal@18:             val = random.randint(0, 100)
bgneal@18:             tree.insert(val)
bgneal@17: 
bgneal@18:         for n in tree:
bgneal@18:             print n.key,
bgneal@18:         print
bgneal@18: 
bgneal@18:         tree.validate()