This repository has been archived on 2021-10-31. You can view files and clone it, but cannot push or open issues or pull requests.
DSA/redblack.py

264 lines
5.3 KiB
Python
Raw Normal View History

2019-05-09 08:15:54 +00:00
#!/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)