bgneal@15: """A red-black tree for Section 3.4, Exercise 4 in Allen Downey's _Think
bgneal@15: Complexity_ book.
bgneal@15: 
bgneal@15: http://greenteapress.com/complexity
bgneal@15: 
bgneal@17: This code is based on the description of the red-black tree at
bgneal@17: http://en.wikipedia.org/wiki/Red-black_tree.
bgneal@17: 
bgneal@17: Some code and ideas were taken from code by Darren Hart at
bgneal@17: http://dvhart.com/darren/files/rbtree.py
bgneal@17: 
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@15: """
bgneal@15: 
bgneal@16: BLACK, RED = range(2)
bgneal@15: 
bgneal@15: class Node(object):
bgneal@15:     """A node class for red-black trees.
bgneal@15: 
bgneal@15:     A node has an optional parent, and optional left and right children. Each
bgneal@15:     node also has a color, either red or black. A node has a key and an optional
bgneal@15:     value. The key is used to order the red black tree by calling the "<"
bgneal@16:     operator when comparing keys. The optional value is useful for using the
bgneal@16:     red-black tree to implement a map datastructure.
bgneal@15: 
bgneal@15:     In a red-black tree, nil children are always considered black.
bgneal@15: 
bgneal@15:     """
bgneal@15: 
bgneal@15:     def __init__(self, key, color=RED, value=None):
bgneal@15:         self.key = key
bgneal@15:         self.value = value
bgneal@15:         self.color = color
bgneal@15: 
bgneal@15:         self.parent = None
bgneal@15:         self.left = None
bgneal@15:         self.right = None
bgneal@15: 
bgneal@15:     @property
bgneal@15:     def grandparent(self):
bgneal@15:         """Return the grandparent of this node, or None if it does not exist."""
bgneal@15:         return self.parent.parent if self.parent else None
bgneal@15: 
bgneal@15:     @property
bgneal@15:     def uncle(self):
bgneal@15:         """Return this node's uncle if it exists, or None if not.
bgneal@15: 
bgneal@15:         An uncle is a parent's sibling.
bgneal@15: 
bgneal@15:         """
bgneal@15:         g = self.grandparent
bgneal@15:         if g:
bgneal@15:             return g.left if g.right is self.parent else g.right
bgneal@15:         return 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@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@17:         if node.left:
bgneal@17:             for n in self._inorder(node.left):
bgneal@17:                 yield n
bgneal@17: 
bgneal@17:         yield node
bgneal@17: 
bgneal@17:         if node.right:
bgneal@17:             for n in self._inorder(node.right):
bgneal@17:                 yield n
bgneal@17: 
bgneal@17:     def insert(self, key, value=None):
bgneal@17:         node = Node(key=key, value=value, color=RED)
bgneal@17: 
bgneal@17:         # trivial case of inserting into an empty tree:
bgneal@17:         if self.root is None:
bgneal@17:             self.root = node
bgneal@17:             self.root.color = BLACK
bgneal@17:             return
bgneal@17: 
bgneal@17:         # Find a spot to insert the new red node:
bgneal@17: 
bgneal@17:         x = self.root
bgneal@17:         while x is not None:
bgneal@17:             p = x
bgneal@17:             if key < x.key:
bgneal@17:                 x = x.left
bgneal@17:             else:
bgneal@17:                 x = x.right
bgneal@17: 
bgneal@17:         node.parent = p     # p is the new node's parent
bgneal@17: 
bgneal@17:         # p now has 1 child; decide if the new node goes on the left or right
bgneal@17: 
bgneal@17:         if key < p.key:
bgneal@17:             p.left = node
bgneal@17:         else:
bgneal@17:             p.right = node
bgneal@17: 
bgneal@17:         # ensure the new tree follows the red-black rules:
bgneal@17: 
bgneal@17:         while True:
bgneal@17:             p = node.parent
bgneal@17: 
bgneal@17:             # Case 1: root node
bgneal@17:             if p is None:
bgneal@17:                 node.color = BLACK
bgneal@17:                 break
bgneal@17: 
bgneal@17:             # Case 2: parent is black
bgneal@17:             if p.color == BLACK:
bgneal@17:                 break
bgneal@17: 
bgneal@17:             # Case 3: parent and uncle are red
bgneal@17:             u = node.uncle
bgneal@17:             if u and u.color == RED:
bgneal@17:                 # repaint parent & uncle black; grandparent becomes red:
bgneal@17:                 p.color = BLACK
bgneal@17:                 u.color = BLACK
bgneal@17:                 gp = node.grandparent
bgneal@17:                 gp.color = RED
bgneal@17: 
bgneal@17:                 # enforce the rules at the grandparent level
bgneal@17:                 node = gp
bgneal@17:                 continue
bgneal@17: 
bgneal@17:             # gp is the grandparent; it can't be None because we know the parent
bgneal@17:             # is red at this point
bgneal@17:             gp = p.parent
bgneal@17: 
bgneal@17:             # Case 4: At this point we know the parent is red and the uncle is
bgneal@17:             # black. 
bgneal@17:             # If the new node is a right child of the parent, and the
bgneal@17:             # parent is on the left of the grandparent => left rotation on the
bgneal@17:             # parent.
bgneal@17:             # If the new node is a left child of the parent, and the parent is
bgneal@17:             # on the right of the grandparent => right rotation on the parent.
bgneal@17: 
bgneal@17:             if node is p.right and p is gp.left:
bgneal@17:                 self._rotate_left(p)
bgneal@17:                 node = node.left
bgneal@17:             elif node is p.left and p is gp.right:
bgneal@17:                 self._rotate_right(p)
bgneal@17:                 node = node.right
bgneal@17: 
bgneal@17:             # Note that case 4 must fall through to case 5 to fix up the former parent
bgneal@17:             # node, which is now the child of the new red node.
bgneal@17: 
bgneal@17:             # Case 5: The parent is red and the uncle is black.
bgneal@17:             # If the new node is the left child of the parent and the parent is
bgneal@17:             # the left child of the grandparent => rotate right on the
bgneal@17:             # grandparent.
bgneal@17:             # If the new node is the right child of the parent and the parent 
bgneal@17:             # is the right child of the grandparent => rotate left on the
bgneal@17:             # grandparent.
bgneal@17: 
bgneal@17:             p = node.parent
bgneal@17:             gp = p.parent
bgneal@17: 
bgneal@17:             p.color = BLACK
bgneal@17:             gp.color = RED
bgneal@17: 
bgneal@17:             if node is p.left and p is gp.left:
bgneal@17:                 self._rotate_right(gp)
bgneal@17:             elif node is p.right and p is gp.right: #TODO: can this be an else?
bgneal@17:                 self._rotate_left(gp)
bgneal@17:             break
bgneal@17: 
bgneal@17:     def _rotate_right(self, node):
bgneal@17:         """Rotate the tree right on the given node."""
bgneal@17: 
bgneal@17:         p = node.parent
bgneal@17:         left = node.left
bgneal@17: 
bgneal@17:         if p:
bgneal@17:             if node is p.right:
bgneal@17:                 p.right = left
bgneal@17:             else:
bgneal@17:                 p.left = left
bgneal@17: 
bgneal@17:         left.parent = p
bgneal@17:         node.left = left.right
bgneal@17: 
bgneal@17:         if node.left:
bgneal@17:             node.left.parent = node
bgneal@17: 
bgneal@17:         left.right = node
bgneal@17:         node.parent = left
bgneal@17: 
bgneal@17:         # fix up the root if necessary
bgneal@17:         if self.root is node:
bgneal@17:             self.root = left
bgneal@17: 
bgneal@17:     def _rotate_left(self, node):
bgneal@17:         """Rotate the tree left on the given node."""
bgneal@17: 
bgneal@17:         p = node.parent
bgneal@17:         right = node.right
bgneal@17: 
bgneal@17:         if p:
bgneal@17:             if node is p.right:
bgneal@17:                 p.right = right
bgneal@17:             else:
bgneal@17:                 p.left = right
bgneal@17: 
bgneal@17:         right.parent = p
bgneal@17:         node.right = right.left
bgneal@17: 
bgneal@17:         if node.right:
bgneal@17:             node.right.parent = node
bgneal@17: 
bgneal@17:         right.left = node
bgneal@17:         node.parent = right
bgneal@17: 
bgneal@17:         # fix up the root if necessary
bgneal@17:         if self.root is node:
bgneal@17:             self.root = right
bgneal@17: 
bgneal@17: 
bgneal@17: if __name__ == '__main__':
bgneal@17:     import random
bgneal@17: 
bgneal@17:     tree = Tree()
bgneal@17: 
bgneal@17:     for i in range(20):
bgneal@17:         val = random.randint(0, 100)
bgneal@17:         tree.insert(val)
bgneal@17: 
bgneal@17:     for n in tree:
bgneal@17:         print n.key,