DSA/tree.py

#!/usr/bin/env python3
# vim: set ts=2 sw=2 et tw=80:

import sys

class Node:
  def __init__(self, k):
    self.key = k
    self.left = None
    self.right = None
    self.parent = None

  def set_left(self, kNode):
    kNode.parent = self
    self.left = kNode

  def set_right(self, kNode):
    kNode.parent = self
    self.right = kNode


tree = Node(77)
tree.set_left(Node(30))
tree.left.set_left(Node(21))
tree.left.set_right(Node(33))
tree.left.right.set_left(Node(31))
tree.left.right.set_right(Node(50))
tree.set_right(Node(78))
tree.right.set_right(Node(80))


# Complexity: Theta(n)
def in_order_walk(t):
  if t is not None:
    in_order_walk(t.left)
    print(t.key, end=' ')
    in_order_walk(t.right)


# Complexity: Theta(n)
def reverse_walk(t):
  if t is not None:
    reverse_walk(t.right)
    print(t.key, end=' ')
    reverse_walk(t.left)


# 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


def insert(t, k):
  insert_node(t, Node(k))


def insert_node(t, node):
  if node.key < t.key:
    if t.left is None:
      t.set_left(node)
    else:
      insert_node(t.left, node)
  else:
    if t.right is None:
      t.set_right(node)
    else:
      insert_node(t.right, node)


def right_rotate(x):
  # assume x is not None
  # assume x.left is not None
  t = x.left
  x.left = t.right
  t.right = x
  return t


def left_rotate(x):
  # assume x is not None
  # assume x.right is not None
  t = x.right
  x.right = t.left
  t.left = x
  return t


def root_insert(t, k):
  if t is None:
    return k
  if k.key > t.key:
    t.set_right(root_insert(t.right, k))
    return left_rotate(t)
  else:
    t.set_left(root_insert(t.left, k))
    return right_rotate(t)


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):
  to_delete = search(t, k)
  if to_delete is None:
    return
  elif to_delete.left is None:
    unlink_me(to_delete, to_delete.right)
  elif to_delete.right is None:
    unlink_me(to_delete, to_delete.left)
  else:
    if abs(to_delete.left.key - to_delete.key) < abs(to_delete.right.key -
        to_delete.key):
      ins = to_delete.right
      new_branch = unlink_me(to_delete, to_delete.left)
      insert_node(new_branch, ins)
    else:
      ins = to_delete.left
      new_branch = unlink_me(to_delete, to_delete.right)
      insert_node(new_branch, ins)


###############################################################################
# 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):
    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 = Node(int(args[0]))
  for i in range(1, len(args)):
    if args[i] == '|':
      break
    print_tree(T)
    print("\nInsert " + str(args[i]) + ":")
    insert(T, int(args[i]))
  for i in range(i+1, len(args)):
    print_tree(T)
    print("\nDelete " + str(args[i]) + ":")
    delete(T, int(args[i]))
  print_tree(T)
  #print(tree.left.right.right.key)
  #print(successor(tree.left.right.right).key)
  #print(predecessor(tree.left.right.right).key)