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@19:     def free(self):
bgneal@19:         """Call this function when removing a node from a tree.
bgneal@19: 
bgneal@19:         It updates the links to encourage garbage collection.
bgneal@19: 
bgneal@19:         """
bgneal@19:         self.link[0] = self.link[1] = None
bgneal@19: 
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@19:         self._size = 0
bgneal@19: 
bgneal@19:     def __len__(self):
bgneal@19:         """To support the len() function by returning the number of elements in
bgneal@19:         the tree.
bgneal@19: 
bgneal@19:         """
bgneal@19:         return self._size
bgneal@17: 
bgneal@17:     def __iter__(self):
bgneal@19:         """Return an iterator to perform an inorder traversal of the tree."""
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@19:         if node:
bgneal@19:             if node.link[0]:
bgneal@19:                 for n in self._inorder(node.link[0]):
bgneal@19:                     yield n
bgneal@17: 
bgneal@19:             yield node
bgneal@17: 
bgneal@19:             if node.link[1]:
bgneal@19:                 for n in self._inorder(node.link[1]):
bgneal@19:                     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@19:         """Insert the (key, value) pair into the tree. Duplicate keys are not
bgneal@19:         allowed in the tree.
bgneal@19: 
bgneal@19:         Returns a tuple of the form (node, flag) where node is either the newly
bgneal@19:         inserted tree node or an already existing node that has the same key.
bgneal@19:         The flag member is a Boolean that will be True if the node was inserted
bgneal@19:         and False if a node with the given key already exists in the tree.
bgneal@19: 
bgneal@19:         """
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@19:             self._size = 1
bgneal@19:             return self.root, True
bgneal@17: 
bgneal@18:         # False/dummy tree root
bgneal@18:         head = Node(key=None, value=None, color=BLACK)
bgneal@19:         d = last = LEFT     # direction variables
bgneal@17: 
bgneal@19:         # Set up helpers
bgneal@18:         t = head
bgneal@18:         g = p = None
bgneal@19:         q = t.link[1] = self.root
bgneal@19: 
bgneal@19:         # Return values
bgneal@19:         target, insert_flag = (None, False)
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@19:                 self._size += 1
bgneal@19:                 insert_flag = True
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@19:                 target = q
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@19:         return target, insert_flag
bgneal@19: 
bgneal@19:     def remove(self, key):
bgneal@19:         """Remove the given key from the tree.
bgneal@19: 
bgneal@19:         Returns True if the key was found and removed and False if the key was
bgneal@19:         not found in the tree.
bgneal@19: 
bgneal@19:         """
bgneal@19:         remove_flag = False     # return value
bgneal@19: 
bgneal@19:         if self.root is not None:
bgneal@19:             # False/dummy tree root
bgneal@19:             head = Node(key=None, value=None, color=BLACK)
bgneal@19:             f = None    # found item
bgneal@19:             d = RIGHT   # direction
bgneal@19: 
bgneal@19:             # Set up helpers
bgneal@19:             q = head
bgneal@19:             g = p = None
bgneal@19:             q.link[d] = self.root
bgneal@19: 
bgneal@19:             # Search and push a red down to fix red violations as we go
bgneal@19:             while q.link[d]:
bgneal@19:                 last = d
bgneal@19: 
bgneal@19:                 # Move the helpers down
bgneal@19:                 g = p
bgneal@19:                 p = q
bgneal@19:                 q = q.link[d]
bgneal@19:                 d = q.key < key
bgneal@19: 
bgneal@19:                 # Save the node with the matching data and keep going; we'll do
bgneal@19:                 # removal tasks at the end
bgneal@19:                 if q.key == key:
bgneal@19:                     f = q
bgneal@19: 
bgneal@19:                 # Push the red node down with rotations and color flips
bgneal@19:                 if not is_red(q) and not is_red(q.link[d]):
bgneal@19:                     if is_red(q.link[not d]):
bgneal@19:                         p.link[last] = self._single_rotate(q, d)
bgneal@19:                         p = p.link[last]
bgneal@19:                     elif not is_red(q.link[not d]):
bgneal@19:                         s = p.link[not last]
bgneal@19: 
bgneal@19:                         if s:
bgneal@19:                             if not is_red(s.link[not last]) and not is_red(s.link[last]):
bgneal@19:                                 # Color flip
bgneal@19:                                 p.color = BLACK
bgneal@19:                                 s.color = RED
bgneal@19:                                 q.color = RED
bgneal@19:                             else:
bgneal@19:                                 d2 = g.link[1] is p
bgneal@19: 
bgneal@19:                                 if is_red(s.link[last]):
bgneal@19:                                     g.link[d2] = self._double_rotate(p, last)
bgneal@19:                                 elif is_red(s.link[not last]):
bgneal@19:                                     g.link[d2] = self._single_rotate(p, last)
bgneal@19: 
bgneal@19:                                 # Ensure correct coloring
bgneal@19:                                 q.color = g.link[d2].color = RED
bgneal@19:                                 g.link[d2].link[0].color = BLACK
bgneal@19:                                 g.link[d2].link[1].color = BLACK
bgneal@19: 
bgneal@19:             # Replace and remove the saved node
bgneal@19:             if f:
bgneal@19:                 f.key, f.value = q.key, q.value
bgneal@19:                 p.link[p.link[1] is q] = q.link[q.link[0] is None]
bgneal@19:                 q.free()
bgneal@19:                 self._size -= 1
bgneal@19:                 remove_flag = True
bgneal@19: 
bgneal@19:             # Update root and make it black
bgneal@19:             self.root = head.link[1]
bgneal@19:             if self.root:
bgneal@19:                 self.root.color = BLACK
bgneal@19: 
bgneal@19:         return remove_flag
bgneal@19: 
bgneal@20:     def find(self, key):
bgneal@20:         """Looks up the key in the tree and returns the corresponding value, or
bgneal@20:         raises a KeyError if it does not exist in the tree.
bgneal@20: 
bgneal@20:         """
bgneal@20:         p = self.root
bgneal@20:         while p:
bgneal@20:             if p.key == key:
bgneal@20:                 return p.value
bgneal@20:             else:
bgneal@20:                 d = p.key < key
bgneal@20:                 p = p.link[d]
bgneal@20: 
bgneal@20:         raise KeyError
bgneal@20: 
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@19:         vals = []
bgneal@19:         vals_set = set()
bgneal@18:         for i in range(25):
bgneal@18:             val = random.randint(0, 100)
bgneal@19:             while val in vals_set:
bgneal@19:                 val = random.randint(0, 100)
bgneal@19:             vals.append(val)
bgneal@19:             vals_set.add(val)
bgneal@19: 
bgneal@18:             tree.insert(val)
bgneal@19:             tree.validate()
bgneal@17: 
bgneal@19:         print 'Inserted in this order:', vals
bgneal@19:         assert len(vals) == len(vals_set)
bgneal@19:         keys = []
bgneal@18:         for n in tree:
bgneal@18:             print n.key,
bgneal@19:             keys.append(n.key)
bgneal@19:         print ' - len:', len(tree)
bgneal@18: 
bgneal@19:         # delete in a random order
bgneal@19:         random.shuffle(keys)
bgneal@19: 
bgneal@19:         for k in keys:
bgneal@19:             print 'Removing', k
bgneal@19:             tree.remove(k)
bgneal@19:             for n in tree:
bgneal@19:                 print n.key,
bgneal@19:             print ' - len:', len(tree)
bgneal@19:             tree.validate()