# HG changeset patch # User Brian Neal # Date 1356573557 21600 # Node ID 92e2879e2e33f483e2102a5843c12f60fa3d895f # Parent 977628018b4bdfbb63ea2a279218d82df9120b1f Rework the red-black tree based on Julienne Walker's tutorial. Insertion is implemented now. Deletion will come next. diff -r 977628018b4b -r 92e2879e2e33 redblacktree.py --- 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()