class rbnode(object):
"""
A node in a red black tree. See Cormen, Leiserson, Rivest, Stein 2nd edition pg 273.
"""
def __init__(self, key):
"Construct."
self._key = key
self._red = False
self._left = None
self._right = None
self._p = None
key = property(fget=lambda self: self._key, doc="The node's key")
red = property(fget=lambda self: self._red, doc="Is the node red?")
left = property(fget=lambda self: self._left, doc="The node's left child")
right = property(fget=lambda self: self._right, doc="The node's right child")
p = property(fget=lambda self: self._p, doc="The node's parent")
def __str__(self):
"String representation."
return str(self.key)
def __repr__(self):
"String representation."
return str(self.key)
class rbtree(object):
"""
A red black tree. See Cormen, Leiserson, Rivest, Stein 2nd edition pg 273.
"""
def __init__(self, create_node=rbnode):
"Construct."
self._nil = create_node(key=None)
"Our nil node, used for all leaves."
self._root = self.nil
"The root of the tree."
self._create_node = create_node
"A callable that creates a node."
root = property(fget=lambda self: self._root, doc="The tree's root node")
nil = property(fget=lambda self: self._nil, doc="The tree's nil node")
def search(self, key, x=None):
"""
Search the subtree rooted at x (or the root if not given) iteratively for the key.
@return: self.nil if it cannot find it.
"""
if None == x:
x = self.root
while x != self.nil and key != x.key:
if key < x.key:
x = x.left
else:
x = x.right
return x
def minimum(self, x=None):
"""
@return: The minimum value in the subtree rooted at x.
"""
if None == x:
x = self.root
while x.left != self.nil:
x = x.left
return x
def maximum(self, x=None):
"""
@return: The maximum value in the subtree rooted at x.
"""
if None == x:
x = self.root
while x.right != self.nil:
x = x.right
return x
def insert_key(self, key):
"Insert the key into the tree."
self.insert_node(self._create_node(key=key))
def insert_node(self, z):
"Insert node z into the tree."
y = self.nil
x = self.root
while x != self.nil:
y = x
if z.key < x.key:
x = x.left
else:
x = x.right
z._p = y
if y == self.nil:
self._root = z
elif z.key < y.key:
y._left = z
else:
y._right = z
z._left = self.nil
z._right = self.nil
z._red = True
self._insert_fixup(z)
def _insert_fixup(self, z):
"Restore red-black properties after insert."
while z.p.red:
if z.p == z.p.p.left:
y = z.p.p.right
if y.red:
z.p._red = False
y._red = False
z.p.p._red = True
z = z.p.p
else:
if z == z.p.right:
z = z.p
self._left_rotate(z)
z.p._red = False
z.p.p._red = True
self._right_rotate(z.p.p)
else:
y = z.p.p.left
if y.red:
z.p._red = False
y._red = False
z.p.p._red = True
z = z.p.p
else:
if z == z.p.left:
z = z.p
self._right_rotate(z)
z.p._red = False
z.p.p._red = True
self._left_rotate(z.p.p)
self.root._red = False
def _left_rotate(self, x):
"Left rotate x."
y = x.right
x._right = y.left
if y.left != self.nil:
y.left._p = x
y._p = x.p
if x.p == self.nil:
self._root = y
elif x == x.p.left:
x.p._left = y
else:
x.p._right = y
y._left = x
x._p = y
def _right_rotate(self, y):
"Left rotate y."
x = y.left
y._left = x.right
if x.right != self.nil:
x.right._p = y
x._p = y.p
if y.p == self.nil:
self._root = x
elif y == y.p.right:
y.p._right = x
else:
y.p._left = x
x._right = y
y._p = x
def check_invariants(self):
"@return: True iff satisfies all criteria to be red-black tree."
def is_red_black_node(node):
"@return: num_black"
# check has _left and _right or neither
if (node.left and not node.right) or (node.right and not node.left):
return 0, False
# check leaves are black
if not node.left and not node.right and node.red:
return 0, False
# if node is red, check children are black
if node.red and node.left and node.right:
if node.left.red or node.right.red:
return 0, False
# descend tree and check black counts are balanced
if node.left and node.right:
# check children's parents are correct
if self.nil != node.left and node != node.left.p:
return 0, False
if self.nil != node.right and node != node.right.p:
return 0, False
# check children are ok
left_counts, left_ok = is_red_black_node(node.left)
if not left_ok:
return 0, False
right_counts, right_ok = is_red_black_node(node.right)
if not right_ok:
return 0, False
# check children's counts are ok
if left_counts != right_counts:
return 0, False
return left_counts, True
else:
return 0, True
num_black, is_ok = is_red_black_node(self.root)
return is_ok and not self.root._red
def write_tree_as_dot(t, f, show_nil=False):
"Write the tree in the dot language format to f."
def node_id(node):
return 'N%d' % id(node)
def node_color(node):
if node.red:
return "red"
else:
return "black"
def visit_node(node):
"Visit a node."
print >> f, " %s [label=\"%s\", color=\"%s\"];" % (node_id(node), node, node_color(node))
if node.left:
if node.left != t.nil or show_nil:
visit_node(node.left)
print >> f, " %s -> %s ;" % (node_id(node), node_id(node.left))
if node.right:
if node.right != t.nil or show_nil:
visit_node(node.right)
print >> f, " %s -> %s ;" % (node_id(node), node_id(node.right))
print >> f, "// Created by rbtree.write_dot()"
print >> f, "digraph red_black_tree {"
visit_node(t.root)
print >> f, "}"
def test_tree(t, keys):
"Insert keys one by one checking invariants and membership as we go."
assert t.check_invariants()
for i, key in enumerate(keys):
for key2 in keys[:i]:
assert t.nil != t.search(key2)
for key2 in keys[i:]:
assert (t.nil == t.search(key2)) ^ (key2 in keys[:i])
t.insert_key(key)
assert t.check_invariants()
if '__main__' == __name__:
import os, sys, numpy.random as R
def write_tree(t, filename):
"Write the tree as an SVG file."
f = open('%s.dot' % filename, 'w')
write_tree_as_dot(t, f)
f.close()
os.system('dot %s.dot -Tsvg -o %s.svg' % (filename, filename))
# test the rbtree
R.seed(2)
size=50
keys = R.randint(-50, 50, size=size)
t = rbtree()
test_tree(t, keys)
write_tree(t, 'tree')
