Added red-black tree implementation

redblack.py

@@ -0,0 +1,263 @@
+#!/usr/bin/env python3
+# vim: set ts=2 sw=2 et tw=80:
+
+import sys
+
+class Tree:
+  def __init__(self, root):
+    self.root = root
+    root.parent = self
+
+  def set_root(self, root):
+    self.root = root
+    root.parent = self
+
+
+class Node:
+  def __init__(self, k):
+    self.key = k
+    self.isBlack = True
+    self.left = None
+    self.right = None
+    self.parent = None
+
+  def set_left(self, kNode):
+    if kNode is not None:
+      kNode.parent = self
+    self.left = kNode
+
+  def set_right(self, kNode):
+    if kNode is not None:
+      kNode.parent = self
+    self.right = kNode
+
+
+def is_black(node):
+  return node is None or node.isBlack
+
+
+def insert(tree, node):
+  y = None
+  x = tree
+  # Imperatively find place to insert node
+  while x is not None:
+    y = x
+    if node.key < x.key:
+      x = x.left
+    else:
+      x = x.right
+
+  node.parent = y
+  if y is None:
+    tree = node
+  elif node.key < y.key:
+    y.left = node
+  else:
+    y.right = node
+  node.isBlack = False
+  insert_fixup(tree, node)
+
+
+def sibling(node):
+  if node.parent.left is node:
+    return node.parent.right
+  else:
+    return node.parent.left
+
+
+def uncle(node):
+  return sibling(node.parent)
+
+
+def right_rotate(x):
+  # assume x is not None
+  # assume x.left is not None
+  # assume not root (x.parent is not None
+  p = x.parent
+  t = x.left
+  x.set_left(t.right)
+  t.set_right(x)
+  if isinstance(p, Tree):
+    p.set_root(t)
+  elif p.left is x:
+    p.set_left(t)
+  else:
+    p.set_right(t)
+
+
+def left_rotate(x):
+  # assume x is not None
+  # assume x.right is not None
+  # assume not root (x.parent is not None)
+  p = x.parent
+  t = x.right
+  x.set_right(t.left)
+  t.set_left(x)
+  if isinstance(p, Tree):
+    p.set_root(t)
+  elif p.left is x:
+    p.set_left(t)
+  else:
+    p.set_right(t)
+
+
+def insert_fixup(tree, node):
+  if isinstance(node.parent, Tree): # if root
+    node.isBlack = True
+  elif is_black(node.parent):
+    # no fixup needed
+    pass
+  elif isinstance(node.parent.parent, Tree):
+    node.parent.isBlack = True
+  elif not is_black(uncle(node)):
+    node.parent.parent.isBlack = False
+    node.parent.isBlack = True
+    if sibling(node.parent) is not None:
+      sibling(node.parent).isBlack = True
+    insert_fixup(tree, node.parent)
+  else:
+    if node.parent.parent.left is node.parent:
+      if node.parent.right is node:
+        left_rotate(node.parent)
+        node = node.left
+      right_rotate(node.parent.parent)
+    else:
+      if node.parent.left is node:
+        right_rotate(node.parent)
+        node = node.right
+      left_rotate(node.parent.parent)
+    node.parent.isBlack = True
+    if sibling(node) is not None:
+      sibling(node).isBlack = False
+    insert_fixup(tree, node.parent)
+
+
+# Complexity (worst): Theta(n)
+def search(tree, k):
+  if tree is None:
+    return None
+  elif tree.key == k:
+    return tree
+  elif k < tree.key:
+    return search(tree.left, k)
+  else:
+    return search(tree.right, k)
+
+
+# Complexity (worst): Theta(n)
+def min(t):
+  if t is None:
+    return None
+  while t.left is not None:
+    t = t.left
+  return t
+
+
+# Complexity (worst): Theta(n)
+def max(t):
+  if t is None:
+    return None
+  while t.right is not None:
+    t = t.right
+  return t
+
+
+def successor(t):
+  if t.right is not None:
+    return min(t.right)
+  while t.parent is not None:
+    if t.parent.left == t:
+      return t.parent
+    else:
+      t = t.parent
+  return None
+
+
+def predecessor(t):
+  if t.left is not None:
+    return max(t.left)
+  while t is not None:
+    if t.parent.right == t:
+      return t.parent
+    else:
+      t = t.parent
+  return None
+
+###############################################################################
+# Code for printing trees, ignore this
+
+
+class Canvas:
+  def __init__(self, width):
+    self.line_width = width
+    self.canvas = []
+
+  def put_char(self, x, y, c):
+    if x < self.line_width:
+      pos = y * self.line_width + x
+      l = len(self.canvas)
+      if pos < l:
+        self.canvas[pos] = c
+      else:
+        self.canvas[l:] = [' '] * (pos - l)
+        self.canvas.append(c)
+
+  def print_out(self):
+    i = 0
+    for c in self.canvas:
+      sys.stdout.write(c)
+      i = i + 1
+      if i % self.line_width == 0:
+        sys.stdout.write('\n')
+    if i % self.line_width != 0:
+      sys.stdout.write('\n')
+
+
+def print_binary_r(t, x, y, canvas):
+  max_y = y
+  if t.left is not None:
+    x, max_y, lx, rx = print_binary_r(t.left, x, y + 2, canvas)
+    x = x + 1
+    for i in range(rx, x):
+      canvas.put_char(i, y + 1, '/')
+
+  middle_l = x
+  for c in (str(t.key) + ("B" if t.isBlack else "R")):
+    canvas.put_char(x, y, c)
+    x = x + 1
+  middle_r = x
+
+  if t.right is not None:
+    canvas.put_char(x, y + 1, '\\')
+    x = x + 1
+    x0, max_y2, lx, rx = print_binary_r(t.right, x, y + 2, canvas)
+    if max_y2 > max_y:
+      max_y = max_y2
+    for i in range(x, lx):
+      canvas.put_char(i, y + 1, '\\')
+    x = x0
+
+  return (x, max_y, middle_l, middle_r)
+
+
+def print_tree(t):
+  print_w(t, 80)
+
+
+def print_w(t, width):
+  canvas = Canvas(width)
+  print_binary_r(t, 0, 0, canvas)
+  canvas.print_out()
+
+###############################################################################
+
+
+if __name__ == "__main__":
+  args = [x for x in sys.argv[1:]]
+  T = Tree(Node(int(args[0])))
+  for i in range(1, len(args)):
+    print_tree(T.root)
+    print("\nInsert " + str(args[i]) + ":")
+    insert(T.root, Node(int(args[i])))
+  print_tree(T.root)
+

tree.py

@@ -134,8 +134,8 @@ def left_rotate(x):
 
 def root_insert(t, k):
   if t is None:
-    return Node(k)
+    return k
-  if k > t.key:
+  if k.key > t.key:
     t.set_right(root_insert(t.right, k))
     return left_rotate(t)
   else:
@@ -145,13 +145,17 @@ def root_insert(t, k):
 
 def unlink_me(node, to_link):
   if node.parent == None:
+    tr = node
     node.key = to_link.key
     node.left = to_link.left
     node.right = to_link.right
+    return node
   elif node.parent.left == node:
     node.parent.left = to_link
+    return to_link
   else:
     node.parent.right = to_link
+    return to_link
 
 
 def delete(t, k):
@@ -165,11 +169,13 @@ def delete(t, k):
   else:
     if abs(to_delete.left.key - to_delete.key) < abs(to_delete.right.key -
         to_delete.key):
-      unlink_me(to_delete, to_delete.left)
+      ins = to_delete.right
-      insert_node(to_delete.left, to_delete.right)
+      new_branch = unlink_me(to_delete, to_delete.left)
+      insert_node(new_branch, ins)
     else:
-      unlink_me(to_delete, to_delete.right)
+      ins = to_delete.left
-      insert_node(to_delete.right, to_delete.left)
+      new_branch = unlink_me(to_delete, to_delete.right)
+      insert_node(new_branch, ins)
 
 
 ###############################################################################