# coding=utf-8
# 红黑树Python实现
# 颜色常量
RED = 0
BLACK = 1
def left_rotate(tree, node):
if not node.right:
return False
node_right = node.right
node_right.p = node.p
if not node.p:
tree.root = node_right
elif node == node.p.left:
node.p.left = node_right
else:
node.p.right = node_right
if node_right.left:
node_right.left.p = node
node.right = node_right.left
node.p = node_right
node_right.left = node
def right_rotate(tree, node):
if not node.left:
return False
node_left = node.left
node_left.p = node.p
if not node.p:
tree.root = node_left
elif node == node.p.left:
node.p.left = node_left
elif node == node.p.right:
node.p.right = node_left
if node_left.right:
node_left.right.p = node
node.left = node_left.right
node.p = node_left
node_left.right = node
def transplant(tree, node_u, node_v):
"""
用 v 替换 u
:param tree: 树的根节点
:param node_u: 将被替换的节点
:param node_v: 替换后的节点
:return: None
"""
if not node_u.p:
tree.root = node_v
elif node_u == node_u.p.left:
node_u.p.left = node_v
elif node_u == node_u.p.right:
node_u.p.right = node_v
# 加一下为空的判断
if node_v:
node_v.p = node_u.p
def tree_maximum(node):
"""
找到以 node 节点为根节点的树的最大值节点 并返回
:param node: 以该节点为根节点的树
:return: 最大值节点
"""
temp_node = node
while temp_node.right:
temp_node = temp_node.right
return temp_node
def tree_minimum(node):
"""
找到以 node 节点为根节点的树的最小值节点 并返回
:param node: 以该节点为根节点的树
:return: 最小值节点
"""
temp_node = node
while temp_node.left:
temp_node = temp_node.left
return temp_node
def preorder_tree_walk(node):
if node:
print (node.value, node.color)
preorder_tree_walk(node.left)
preorder_tree_walk(node.right)
class RedBlackTreeNode(object):
def __init__(self, value):
self.value = value
self.left = None
self.right = None
self.p = None
self.color = RED
class RedBlackTree(object):
def __init__(self):
self.root = None
def insert(self, node):
# 找到最接近的节点
temp_root = self.root
temp_node = None
while temp_root:
temp_node = temp_root
if node.value == temp_node.value:
return False
elif node.value > temp_node.value:
temp_root = temp_root.right
else:
temp_root = temp_root.left
# 在相应位置插入节点
if not temp_node:
self.root = node
node.color = BLACK
elif node.value < temp_node.value:
temp_node.left = node
node.p = temp_node
else:
temp_node.right = node
node.p = temp_node
# 调整树
self.insert_fixup(node)
def insert_fixup(self, node):
if node.value == self.root.value:
return
# 为什么是这个终止条件？
# 因为如果不是这个终止条件那就不需要调整
while node.p and node.p.color == RED:
# 只要进入循环则必有祖父节点 否则父节点为根节点 根节点颜色为黑色 不会进入循环
if node.p == node.p.p.left:
node_uncle = node.p.p.right
# 1. 没有叔叔节点 若此节点为父节点的右子 则先左旋再右旋 否则直接右旋
# 2. 有叔叔节点 叔叔节点颜色为黑色
# 3. 有叔叔节点 叔叔节点颜色为红色 父节点颜色置黑 叔叔节点颜色置黑 祖父节点颜色置红 continue
# 注: 1 2 情况可以合为一起讨论 父节点为祖父节点右子情况相同 只需要改指针指向即可
if node_uncle and node_uncle.color == RED:
node.p.color = BLACK
node_uncle.color = BLACK
node.p.p.color = RED
node = node.p.p
continue
elif node == node.p.right:
left_rotate(self, node.p)
node = node.left
node.p.color = BLACK
node.p.p.color = RED
right_rotate(self, node.p.p)
return
elif node.p == node.p.p.right:
node_uncle = node.p.p.left
if node_uncle and node_uncle.color == RED:
node.p.color = BLACK
node_uncle.color = BLACK
node.p.p.color = RED
node = node.p.p
continue
elif node == node.p.left:
right_rotate(self, node)
node = node.right
node.p.color = BLACK
node.p.p.color = RED
left_rotate(self, node.p.p)
return
# 最后记得把根节点的颜色改为黑色 保证红黑树特性
self.root.color = BLACK
def delete(self, node):
# 找到以该节点为根节点的右子树的最小节点
node_color = node.color
if not node.left:
temp_node = node.right
transplant(self, node, node.right)
elif not node.right:
temp_node = node.left
transplant(self, node, node.left)
else:
# 最麻烦的一种情况 既有左子 又有右子 找到右子中最小的做替换 类似于二分查找树的删除
node_min = tree_minimum(node.right)
node_color = node_min.color
temp_node = node_min.right
if node_min.p != node:
transplant(self, node_min, node_min.right)
node_min.right = node.right
node_min.right.p = node_min
transplant(self, node, node_min)
node_min.left = node.left
node_min.left.p = node_min
node_min.color = node.color
# 当删除的节点的颜色为黑色时 需要调整红黑树
if node_color == BLACK:
self.delete_fixup(temp_node)
def delete_fixup(self, node):
# 实现过程还需要理解 比如为什么要删除 为什么是那几种情况
while node != self.root and node.color == BLACK:
if node == node.p.left:
node_brother = node.p.right
if node_brother.color == RED:
node_brother.color = BLACK
node.p.color = RED
left_rotate(self, node.p)
node_brother = node.p.right
if (not node_brother.left or node_brother.left.color == BLACK) and \
(not node_brother.right or node_brother.right.color == BLACK):
node_brother.color = RED
node = node.p
else:
if not node_brother.right or node_brother.right.color == BLACK:
node_brother.color = RED
node_brother.left.color = BLACK
right_rotate(self, node_brother)
node_brother = node.p.right
node_brother.color = node.p.color
node.p.color = BLACK
node_brother.right.color = BLACK
left_rotate(self, node.p)
node = self.root
break
else:
node_brother = node.p.left
if node_brother.color == RED:
node_brother.color = BLACK
node.p.color = RED
left_rotate(self, node.p)
node_brother = node.p.right
if (not node_brother.left or node_brother.left.color == BLACK) and \
(not node_brother.right or node_brother.right.color == BLACK):
node_brother.color = RED
node = node.p
else:
if not node_brother.left or node_brother.left.color == BLACK:
node_brother.color = RED
node_brother.right.color = BLACK
left_rotate(self, node_brother)
node_brother = node.p.left
node_brother.color = node.p.color
node.p.color = BLACK
node_brother.left.color = BLACK
right_rotate(self, node.p)
node = self.root
break
node.color = BLACK
def main():
number_list = (7, 4, 1, 8, 5, 2, 9, 6, 3)
tree = RedBlackTree()
for number in number_list:
node = RedBlackTreeNode(number)
tree.insert(node)
del node
preorder_tree_walk(tree.root)
tree.delete(tree.root)
preorder_tree_walk(tree.root)
if __name__ == '__main__':
main()
教材: 算法导论(第三版)
