DSA/GA3/redblack.py

#!/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 not isinstance(node.parent.parent, Tree) and 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.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

###############################################################################
# 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)