Skip to content

Instantly share code, notes, and snippets.

@DanielNill
Created June 4, 2019 05:03
Show Gist options
  • Select an option

  • Save DanielNill/0c0a5859fd04a70daa3e02ef977339a8 to your computer and use it in GitHub Desktop.

Select an option

Save DanielNill/0c0a5859fd04a70daa3e02ef977339a8 to your computer and use it in GitHub Desktop.
class NilNode(object):
color = "black"
parent = None
left = None
right = None
NIL = NilNode()
class Node(object):
left = NIL
right = NIL
parent = None
color = "red"
def __init__(self, v):
self.v = v
class RBTree(object):
root = NIL
def replace(self, a, b):
if a == self.root:
self.root = b
elif a == a.parent.left:
a.parent.left = b
else:
a.parent.right = b
if b:
b.parent = a.parent
def delete(self, v):
node = self.root
match = None
while node:
if node.v == v:
match = node
break
if v < node.v:
node = node.left
else:
node = node.right
if not match: return False
rep = match
rep_color = rep.color
if match.left == NIL:
child = match.right
self.replace(match, match.right)
elif match.right == NIL:
child = match.left
self.replace(match, match.left)
else:
rep = match.right
while rep.left != NIL:
rep = rep.left
rep_color = rep.color
child = rep.right
if rep.parent == match:
child.parent = rep
else:
self.replace(rep, rep.right)
rep.right = match.right
rep.right.parent = rep
self.replace(match, rep)
rep.left = match.left
rep.left.parent = rep
rep.color = match.color
while child != self.root and rep_color == "black" and child.color == "black":
if child == child.parent.left:
sibling = child.parent.right
if sibling.color == "red":
sibling.color = "black"
child.parent.color = "red"
self.rotate_left(child.parent)
sibling = child.parent.right
if sibling.left.color == "black" and sibling.right.color == "black":
sibling.color = "red"
child = child.parent
else:
if sibling.right.color == "black":
sibling.left.color = "black"
sibling.color = "red"
self.rotate_right(sibling)
sibling = child.parent.right
sibling.color = child.parent.color
child.parent.color = "black"
sibling.right.color = "black"
self.rotate_left(child.parent)
child = self.root
else:
sibling = child.parent.left
if sibling.color == "red":
sibling.color = "black"
child.parent.color = "red"
self.rotate_right(child.parent)
sibling = child.parent.left
if sibling.left.color == "black" and sibling.right.color == "black":
sibling.color = "red"
child = child.parent
else:
if sibling.left.color == "black":
sibling.right.color = "black"
sibling.color = "red"
self.rotate_left(sibling)
sibling = child.parent.left
sibling.color = child.parent.color
child.parent.color = "black"
sibling.left.color = "black"
self.rotate_right(child.parent)
child = self.root
child.color = "black"
def insert(self, v):
node = Node(v)
cur = self.root
prev = None
while cur != NIL:
prev = cur
if v < cur.v:
cur = cur.left
else:
cur = cur.right
if not prev:
self.root = node
elif v < prev.v:
prev.left = node
else:
prev.right = node
node.parent = prev
if not node.parent or not node.parent.parent: return
while node != self.root and node.parent.color == "red":
if node.parent == node.parent.parent.left:
uncle = node.parent.parent.right
if uncle.color == "red":
uncle.color = "black"
node.parent.color = "black"
node.parent.parent.color = "red"
node = node.parent.parent
else:
if node == node.parent.right:
node = node.parent
self.rotate_left(node)
node.parent.color = "black"
node.parent.parent.color = "red"
self.rotate_right(node.parent.parent)
else:
uncle = node.parent.parent.left
if uncle.color == "red":
uncle.color = "black"
node.parent.color = "black"
node.parent.parent.color = "red"
node = node.parent.parent
else:
if node == node.parent.left:
node = node.parent
self.rotate_right(node)
node.parent.color = "black"
node.parent.parent.color = "red"
self.rotate_left(node.parent.parent)
self.root.color = "black"
def rotate_left(self, node):
right_child = node.right
grandchild = right_child.left
right_child.left = node
node.right = grandchild
if node == self.root:
self.root = right_child
elif node.parent.left == node:
node.parent.left = right_child
else:
node.parent.right = right_child
right_child.parent = node.parent
node.parent = right_child
grandchild.parent = node
def rotate_right(self, node):
left_child = node.left
grandchild = left_child.right
left_child.right = node
node.left = grandchild
if node == self.root:
self.root = left_child
elif node.parent.left == node:
node.parent.left = left_child
else:
node.parent.right = left_child
left_child.parent = node.parent
node.parent = left_child
grandchild.parent = node
def preorder(self, n):
if n != NIL:
print(n.v)
self.preorder(n.left)
self.preorder(n.right)
t = RBTree()
t.insert(10)
t.insert(20)
t.insert(30)
t.insert(40)
t.insert(50)
t.insert(60)
t.preorder(t.root)
# 20 10 40 30 50 60
t.delete(40)
print("###")
t.preorder(t.root)
# 20 10 50 30 60
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment