原理讲解参考：https://www.cnblogs.com/skywang12345/p/3245399.html
'''
1) 每个节点是黑色或者红色。
2) 根节点是黑色。
3) 每个叶子节点(nil)是黑色。
4) 红色节点的子节点必须是黑色。
5) 从一个节点到该节点的子孙节点的所有路径上包含相同数目的黑节点。
'''
import random
class Color:
Red = 1
Black = 2
class NullNode:
def __init__(self):
self.color = Color.Black
self.parent = None
self.left = None
self.right = None
class RbTreeNode:
def __init__(self, key, color):
self.key = key
self.color = color
self.parent = NullNode()
self.left = NullNode()
self.right = NullNode()
def format(self):
if self.color == Color.Red:
return " 33[31m%s 33[0m" % self.key
else:
return " 33[34m%s 33[0m" % self.key
def left_rotate(t, x):
p = x.parent
r = x.right
# 右节点设为父节点
r.parent = p
if x == p.left:
p.left = r
else:
p.right = r
# 右左节点设置x的右节点
r.left.parent = x
x.right = r.left
# x设为左节点
r.left = x
x.parent = r
if not r.parent or isinstance(r.parent, NullNode):
t = r
return t
def right_rotate(t, x):
p = x.parent
l = x.left
l.parent = p
if x == p.left:
p.left = l
else:
p.right = l
l.right.parent = x
x.left = l.right
l.right = x
x.parent = l
if not l.parent or isinstance(l.parent, NullNode):
t = l
return t
def insert(t, key):
z = RbTreeNode(key, Color.Red)
y = NullNode()
x = t
while x and not isinstance(x, NullNode):
y = x
if z.key < x.key:
x = x.left
else:
x = x.right
z.parent = y
if not y or isinstance(y, NullNode):
t = z
print(t)
elif z.key < y.key:
y.left = z
else:
y.right = z
print_tree(t)
return insert_fixup(t, z)
# 有bug,待修复
def insert_fixup(t, z):
while z.parent.color == Color.Red:
if z.parent == z.parent.parent.left:
y = z.parent.parent.right
if y.color == Color.Red:
z.parent.color = Color.Black
y.color = Color.Black
z.parent.parent.color = Color.Red
z = z.parent.parent
else:
if z == z.parent.right:
z = z.parent
t = left_rotate(t, z)
z.parent.color = Color.Black
z.parent.parent.color = Color.Red
t = right_rotate(t, z.parent.parent)
else:
y = z.parent.parent.left
if y.color == Color.Red:
z.parent.color = Color.Black
y.color = Color.Black
z.parent.parent.color = Color.Red
z = z.parent.parent
else:
if z == z.parent.left:
z = z.parent
t = right_rotate(t, z)
z.parent.color = Color.Black
z.parent.parent.color = Color.Red
t = left_rotate(t, z.parent.parent)
t.color = Color.Black
return t
def delete(t, key):
z = NullNode()
y = NullNode()
x = t
while x and not isinstance(x, NullNode):
if key < x.key:
x = x.left
elif key > x.key:
x = x.right
else:
z = x
break
if isinstance(z, NullNode):
return t
if isinstance(z.left, NullNode) or isinstance(z.right, NullNode):
y = z
else:
y = tree_successor(z)
if not isinstance(y.left, NullNode):
x = y.left
else:
x = y.right
x.parent = y.parent
if isinstance(y.parent, NullNode):
t = x
elif y == y.parent.left:
y.parent.left = x
else:
y.parent.right = x
if y != z:
z.key = y.key
if y.color == Color.Black:
t = delete_fixup(t, x)
return t
def delete_fixup(t, x):
while x != t and x.color == Color.Black:
if x == x.parent.left:
w = x.parent.right
if w.color == Color.Red:
w.color = Color.Black
t = left_rotate(t, x.parent)
w = x.parent.right
if w.left.color == Color.Black and w.right.color == Color.Black:
w.color = Color.Red
x = x.parent
elif w.right.color == Color.Black:
w.left.color = Color.Black
w.color = Color.Red
t = right_rotate(t, w)
w = x.parent.right
w.color = x.parent.color
x.parent.color = Color.Black
w.right.color = Color.Black
t = left_rotate(t, x.parent)
x = t
else:
w = x.parent.left
if w.color == Color.Red:
w.color = Color.Black
t = right_rotate(t, x.parent)
w = x.parent.left
if w.right.color == Color.Black and w.left.color == Color.Black:
w.color = Color.Red
x = x.parent
elif w.left.color == Color.Black:
w.right.color = Color.Black
w.color = Color.Red
t = left_rotate(t, w)
w = x.parent.left
w.color = x.parent.color
x.parent.color = Color.Black
w.left.color = Color.Black
t = right_rotate(t, x.parent)
x = t
x.color = Color.Black
return t
def tree_successor(x):
y = x.right
z = NullNode()
while y and not isinstance(y, NullNode):
z = y
y = y.left
return z
def set_parent(p, l, r):
p.left = l
p.right = r
l.parent = p
r.parent = p
def format_tree(t):
if not t or isinstance(t, NullNode):
return 0, []
left_indent, left_tree = format_tree(t.left)
right_indent, right_tree = format_tree(t.right)
tree = []
tree.append("."*left_indent + t.format() + "."*right_indent)
n = 0
while n < len(left_tree) or n < len(right_tree):
line = ""
if n < len(left_tree):
line += left_tree[n]
else:
line += " "*left_indent
line += " "*len(str(t.key))
if n < len(right_tree):
line += right_tree[n]
else:
line += " "*right_indent
tree.append(line)
n += 1
indent = left_indent + len(str(t.key)) + right_indent
return indent, tree
def print_tree(t):
# print("------------")
_, tree = format_tree(t)
for line in tree:
print(line)
if __name__ == "__main__":
# a = RbTreeNode("a", Color.Red)
# b = RbTreeNode("b", Color.Red)
# c = RbTreeNode("c", Color.Red)
# d = RbTreeNode("d", Color.Red)
# set_parent(a, b, NullNode())
# set_parent(b, c, d)
# print_tree(a)
# # left_rotate(a, b)
# right_rotate(a, b)
# print_tree(a)
# a = RbTreeNode("a", Color.Black)
# b = RbTreeNode("b", Color.Red)
# set_parent(a, NullNode(), b)
# t = a
# print_tree(t)
# t = insert(t, "c")
# print_tree(t)
t = NullNode()
for i in range(30):
n = random.randint(0, 10)
print("------ ", n)
t = insert(t, n)
print_tree(t)
print("*******************")
for i in range(10):
n = random.randint(0, 10)
print("------ ", n)
t = delete(t, n)
print_tree(t)
# for key in [chr(x) for x in range(ord('a'), ord('z') + 1)]:
# t = insert(t, key)
# print_tree(t)
