Python实现红黑树的插入操作
本专栏中的上一篇文章介绍了什么是红黑树，以及红黑树的旋转和变色。
本文使用Python实现红黑树的插入操作。
先将红黑树的5条特性列出来：
1. 节点是红色或黑色。
2. 根节点是黑色。
3. 所有叶子节点都是黑色的空节点。(叶子节点是NIL节点或NULL节点)
4. 每个红色节点的两个子节点都是黑色节点。(从每个叶子节点到根的所有路径上不能有两个连续的红色节点)
5. 从任一节点到其每个叶子节点的所有路径都包含相同数目的黑色节点。
一、代码准备
红黑树是一种特殊的二叉搜索树，所以先写好二叉搜索树的代码，这部分不详细介绍，后面直接使用。
# coding=utf-8
class RBNode(object):
"""节点类"""
def __init__(self, data, left_child=None, right_child=None, color='red'):
self.data = data
self.parent = None
self.left_child = left_child
self.right_child = right_child
self.color = color
class RBBinaryTree(object):
"""红黑树类"""
def __init__(self):
self.__root = None
self.prefix_branch = '├'
self.prefix_trunk = '|'
self.prefix_leaf = '└'
self.prefix_empty = ''
self.prefix_left = '─L─'
self.prefix_right = '─R─'
def is_empty(self):
return not self.__root
@property
def root(self):
return self.__root
@root.setter
def root(self, value):
self.__root = value if isinstance(value, RBNode) else RBNode(value)
def insert(self, root, value):
"""二叉搜索树插入节点-递归"""
node = value if isinstance(value, RBNode) else RBNode(value)
if self.is_empty():
self.root = node
return
if root is None:
root = node
elif node.data < root.data:
root.left_child = self.insert(root.left_child, value)
root.left_child.parent = root
elif node.data > root.data:
root.right_child = self.insert(root.right_child, value)
root.right_child.parent = root
return root
def search(self, root, data):
"""二叉搜索树的查询操作"""
if root is None:
return
if root.data == data:
return root
elif data < root.data:
return self.search(root.left_child, data)
elif data > root.data:
return self.search(root.right_child, data)
def show_tree(self):
if self.is_empty():
print('空二叉树')
return
print('-' * 20)
print("\033[31m{}\033[0m".format(str(self.root.data))) if self.root.color is 'red' else print(str(self.root.data))
self.__print_tree(self.__root)
print('-' * 20)
def __print_tree(self, node, prefix=None):
if prefix is None:
prefix, prefix_left_child = '', ''
else:
prefix = prefix.replace(self.prefix_branch, self.prefix_trunk).replace(self.prefix_leaf, self.prefix_empty)
prefix_left_child = prefix.replace(self.prefix_leaf, self.prefix_empty)
if self.has_child(node):
if node.right_child is not None:
if node.right_child.color is 'red':
print(prefix + self.prefix_branch + self.prefix_right + "\033[31m{}\033[0m".format(str(node.right_child.data)))
else:
print(prefix + self.prefix_branch + self.prefix_right + str(node.right_child.data))
if self.has_child(node.right_child):
self.__print_tree(node.right_child, prefix + self.prefix_branch + ' ')
else:
print(prefix + self.prefix_branch + self.prefix_right)
if node.left_child is not None:
if node.left_child.color is 'red':
print(prefix + self.prefix_leaf + self.prefix_left + "\033[31m{}\033[0m".format(str(node.left_child.data)))
else:
print(prefix + self.prefix_leaf + self.prefix_left + str(node.left_child.data))
if self.has_child(node.left_child):
prefix_left_child += ' '
self.__print_tree(node.left_child, self.prefix_leaf + prefix_left_child)
else:
print(prefix + self.prefix_leaf + self.prefix_left)
def has_child(self, node):
return node.left_child is not None or node.right_child is not None
1. 定义了一个节点类 RBNode ，用于创建新节点，添加到红黑树中，节点中定义了节点的颜色属性 color 。
color 默认为 red，即默认插入的节点为红节点，这样可以降低破坏红黑树特性的可能性。不管新节点是什么颜色，特性3都不可能被破坏，特性1、2、4都有可能被破坏。如果插入的节点是黑色，则一定会破坏特性5，需要进行调整，如果插入的节点是红色，则一定不会破坏特性5。
2. 定义了红黑树类 RBBinaryTree ，类中实现了按树形结构打印红黑树的方法 show_tree()，并且根据红黑树的节点颜色，打印值时打印对应的颜色。还有二叉搜索树的插入方法 insert(root, value) 和搜索方法 search(root, data) 。
二、实现红黑树的旋转方法
红黑树的旋转分为左旋和右旋。
1. 红黑树的左旋
左旋：以某个节点作为支点(旋转节点)，其右子节点变为旋转节点的父节点，右子节点的左子节点变为旋转节点的右子节点，旋转节点的左子节点保持不变。右子节点的左子节点相当于从右子节点上“断开”，重新连接到旋转节点上。
def left_rotate(self, node):
"""红黑树左旋"""
parent_node, right_node = node.parent, node.right_child
if not right_node:
return
# 1.node是旋转节点，将旋转节点的右子节点的左子节点变为旋转节点的右子节点
node.right_child = right_node.left_child
if node.right_child:
node.right_child.parent = node
# 2.将旋转节点修改为右子节点的左子节点
right_node.left_child, node.parent = node, right_node
# 3.将右子节点替换旋转节点的位置，作为旋转节点父节点的子节点
if not parent_node:
self.root = right_node
else:
if parent_node.left_child == node:
parent_node.left_child = right_node
else:
parent_node.right_child = right_node
right_node.parent = parent_node
left_rotate(node): 以node为旋转节点对红黑树进行左旋操作。
左旋的代码分三步实现：
1. 将旋转节点的右子节点的左子节点变为旋转节点的右子节点。
2. 将旋转节点变为右子节点的左子节点。
3. 将右子节点替换旋转节点的位置，作为旋转节点父节点的子节点。
现在先用二叉搜索树的插入方法生成一棵二叉搜索树。
if __name__ == '__main__':
tree = RBBinaryTree()
data = [50, 77, 55, 29, 10, 30, 66, 18, 80, 51, 90]
for i in data:
tree.insert(tree.root, i)
tree.show_tree()
运行结果：
插入数据后的二叉搜索树结构如下图，因为节点默认都是红色的，所以现在节点全是红节点。这里只是为了展示左旋功能，后面实现红黑树的插入方法后，就可以生成正常的红黑树了。(本文中了为了简洁，结构图中忽略了NIL或NULL节点)
如以节点77作为旋转节点，对这棵二叉搜索树进行左旋。
node = tree.search(tree.root, 77)
tree.left_rotate(node)
tree.show_tree()
运行结果：
左旋后的二叉搜索树结构如下图：
2. 红黑树的右旋
右旋：以某个节点作为支点(旋转节点)，其左子节点变为旋转节点的父节点，左子节点的右子节点变为旋转节点的左子节点，旋转节点的右子节点保持不变。左子节点的右子节点相当于从左子节点上“断开”，重新连接到旋转节点上。
def right_rotate(self, node):
"""红黑树右旋"""
parent_node, left_node = node.parent, node.left_child
if not left_node:
return
# 1.node是旋转节点，将旋转节点的左子节点的右子节点变为旋转节点的左子节点
node.left_child = left_node.right_child
if node.left_child:
node.left_child.parent = node
# 2.将旋转节点修改为左子节点的右子节点
left_node.right_child, node.parent = node, left_node
# 3.将左子节点替换旋转节点的位置，作为旋转节点父节点的子节点
if not parent_node:
self.root = left_node
else:
if parent_node.left_child == node:
parent_node.left_child = left_node
else:
parent_node.right_child = left_node
left_node.parent = parent_node
right_rotate(node): 以node为旋转节点对红黑树进行右旋操作。
右旋的代码分三步实现：
1. 将旋转节点的左子节点的右子节点变为旋转节点的左子节点。
2. 将旋转节点变为左子节点的右子节点。
3. 将左子节点替换旋转节点的位置，作为旋转节点父节点的子节点。
以节点80作为旋转节点，对上一步左旋后的二叉搜索树进行右旋。左旋和右旋是互逆的，右旋后二叉搜索树又回到了最初的结构。
node = tree.search(tree.root, 80)
tree.right_rotate(node)
tree.show_tree()
三、实现红黑树的变色方法
变色：将节点的颜色由红变黑或由黑变红。
def change_color(self, node):
"""红黑树变色"""
if node.color is 'red':
node.color = 'black'
elif node.color is 'black':
node.color = 'red'
change_color(node): 修改节点node的颜色。红变黑，黑变红。
如修改节点10和节点50的颜色。
tree.change_color(tree.search(tree.root, 50))
tree.change_color(tree.search(tree.root, 10))
tree.show_tree()
运行结果：
变色后的二叉搜索树结构如下图。
四、实现红黑树的插入方法
一棵红黑树，一开始是满足5条特性的，插入新节点后，如果特性被破坏了，就要进行调整，使红黑树重新满足5条特性。
可以将这个过程分解成两个大步骤，第一步是按二叉搜索树将节点插入到对应的位置，这里直接调用已经实现的 insert() 方法即可。第二步是判断是否破坏了红黑树的特性，如果破坏了，就使用旋转和变色来对红黑树进行调整。
为了方便理解，先申明几个需要使用到的术语：
因素节点：因为这个节点的变化，使红黑树的结构发生了变化，用F(factor_node)表示。
父节点：因素节点的父节点，用P(parent_node)表示。
祖父节点：因素节点的父节点的父节点，用G(grandparent_node)表示。
叔节点：因素节点的父节点的兄弟节点，用U(uncle_node)表示。
插入新节点后，是否需要对红黑树进行调整，以及怎么调整，分为如下几种情况：
1. 如果红黑树是一棵空树，则将新节点添加到根节点的位置，并将节点的颜色修改为黑色。
2. 如果新节点的父节点是黑节点，则插入新节点后，不需要做调整。
3. 如果新节点的父节点是红节点(不满足特性4)，根据叔节点的颜色，分为两类。
3.1 叔节点为黑色，叔节点为黑色包含了叔节点为空的情况，因为红黑树的空节点(叶子节点)一定是黑色的。
叔节点为黑色时，根据节点的结构关系，可以分为四种情况。
3.1.1 新节点是父节点的左子节点，父节点也是祖父节点的左子节点(左左结构)。将父节点红变黑，祖父节点黑变红，然后以祖父节点作为旋转节点右旋。
这里要注意两点：
(1). 新节点的父节点的另一个子节点一定是叶子节点。因为插入新节点前红黑树是满足5条特性的，假设父节点还有一个非空子节点，如果这个节点红节点，则不满足特性4，如果这个节点是黑节点，则不满足特性5。
(2). 在插入新节点时，如果新节点的父节点是红色，祖父节点是黑色，则叔节点一定不会是非空黑节点，要么是叶子节点要么是红节点。因为插入新节点前红黑树是满足5条特性的，如果叔节点是一个非空黑节点，则红黑树不满足特性5。
在后面的3.2中，叔节点是红节点时，进行第一次调整后，还不一定能调整完成，祖父节点成为新的因素节点，可能会出现因素节点F为红、父节点P为红、祖父节点G为黑、叔节点U为黑的情况。所以这种状态其实是调整过程中产生的中间状态，3.1中先给处理方式，等介绍完叔节点为红节点的情况，就可以连贯到一起了。
3.1.2 新节点是父节点的右子节点，父节点是祖父节点的左子节点(左右结构)。先以父节点作为旋转节点左旋，变成3.1.1的结构，再按3.1.1的方式处理。
3.1.3 新节点是父节点的右子节点，父节点也是祖父节点的右子节点(右右结构)。将父节点红变黑，祖父节点黑变红，然后以祖父节点作为旋转节点左旋。
3.1.4 新节点是父节点的左子节点，父节点是祖父节点的右子节点(右左结构)。先以父节点作为旋转节点右旋，变成3.1.3的结构，再按3.1.3的方式处理。
3.2 叔节点为红色，这种情况将父节点和叔节点红变黑，将祖父节点黑变红。
祖父节点由黑变红后，因为不知道祖父节点的父节点是什么颜色，所以还需要进行判断，将祖父节点作为新的因素节点，递归进入下一次调整，可以分为如下四种情况。
3.2.1 如果祖父节点是根节点(不满足特性1)，要将祖父节点重新变成黑色，并且调整结束。
3.2.2 如果其父节点是黑色，则不需再次调整。
3.2.3 如果其父节点是红色，叔节点是黑色，则符合上面3.1的情况，则递归调用3.1处理。
3.2.4 如果其父节点是红色，叔节点也是红色，则符合当前3.2的情况，递归再进行一次3.2处理。如果处理完成后，新的因素节点还符合3.2的情况，则继续递归，直到退出循环。
根据上面的分析，代码实现如下：
def rb_insert(self, value):
"""红黑树插入"""
node = value if isinstance(value, RBNode) else RBNode(value)
if self.search(self.root, node.data):
return
if self.is_empty():
node.color = 'black'
self.root = node
return
self.insert(self.root, node)
factor_node = node
while True:
parent_node = factor_node.parent
if parent_node.color is 'red':
grandparent_node = parent_node.parent
if parent_node is grandparent_node.left_child:
uncle_node = grandparent_node.right_child
else:
uncle_node = grandparent_node.left_child
# 如果父节点为红节点且叔节点为黑节点
if uncle_node is None or uncle_node and uncle_node.color is 'black':
if parent_node == grandparent_node.left_child:
# 先左旋为左左结果，然后父节点和祖父节点变色，再右旋
if factor_node == parent_node.right_child:
self.left_rotate(parent_node)
self.change_color(factor_node)
else:
self.change_color(parent_node)
self.change_color(grandparent_node)
self.right_rotate(grandparent_node)
elif parent_node == grandparent_node.right_child:
# 先右旋为右右结构，然后父节点和祖父节点变色，再左旋
if factor_node == parent_node.left_child:
self.right_rotate(parent_node)
self.change_color(factor_node)
else:
self.change_color(parent_node)
self.change_color(grandparent_node)
self.left_rotate(grandparent_node)
break
# 如果父节点为红节点且叔节点也为红节点
elif uncle_node and uncle_node.color is 'red':
# 父节点和叔节点变色，祖父节点变色(祖父节点是根节点除外)
self.change_color(parent_node)
self.change_color(uncle_node)
if grandparent_node != self.root:
self.change_color(grandparent_node)
# 祖父节点变成红节点后，将祖父节点作为新的因素节点，判断其父节点，避免不满足特性4
factor_node = grandparent_node
else:
break
rb_insert(value): 红黑树的插入操作，分析了插入操作的所有情况及每种情况的处理方式，代码的实现就相对简单了。
还是一开始的数据，使用红黑树的插入方式依次插入。
if __name__ == '__main__':
tree = RBBinaryTree()
data = [50, 77, 55, 29, 10, 30, 66, 18, 80, 51, 90]
for i in data:
# tree.insert(tree.root, i)
tree.rb_insert(i)
tree.show_tree()
运行结果如下：
得到的红黑树结构如下图：
由于插入操作分的情况比较多，所以在看插入操作的分析时，一定要先理解旋转和变色操作。
实现红黑树的代码后，可以看出，每插入一个新节点，红黑树都是满足5条特性的，而有一些红黑树不一定是一个节点一个节点地添加得到的。
五、完整代码
# coding=utf-8
class RBNode(object):
"""节点类"""
def __init__(self, data, left_child=None, right_child=None, color='red'):
self.data = data
self.parent = None
self.left_child = left_child
self.right_child = right_child
self.color = color
class RBBinaryTree(object):
"""红黑树类"""
def __init__(self):
self.__root = None
self.prefix_branch = '├'
self.prefix_trunk = '|'
self.prefix_leaf = '└'
self.prefix_empty = ''
self.prefix_left = '─L─'
self.prefix_right = '─R─'
def is_empty(self):
return not self.__root
@property
def root(self):
return self.__root
@root.setter
def root(self, value):
self.__root = value if isinstance(value, RBNode) else RBNode(value)
def left_rotate(self, node):
"""红黑树左旋"""
parent_node, right_node = node.parent, node.right_child
if not right_node:
return
# 1.node是旋转节点，将旋转节点的右子节点的左子节点变为旋转节点的右子节点
node.right_child = right_node.left_child
if node.right_child:
node.right_child.parent = node
# 2.将旋转节点修改为右子节点的左子节点
right_node.left_child, node.parent = node, right_node
# 3.将右子节点替换旋转节点的位置，作为旋转节点父节点的子节点
if not parent_node:
self.root = right_node
else:
if parent_node.left_child == node:
parent_node.left_child = right_node
else:
parent_node.right_child = right_node
right_node.parent = parent_node
def right_rotate(self, node):
"""红黑树右旋"""
parent_node, left_node = node.parent, node.left_child
if not left_node:
return
# 1.node是旋转节点，将旋转节点的左子节点的右子节点变为旋转节点的左子节点
node.left_child = left_node.right_child
if node.left_child:
node.left_child.parent = node
# 2.将旋转节点修改为左子节点的右子节点
left_node.right_child, node.parent = node, left_node
# 3.将左子节点替换旋转节点的位置，作为旋转节点父节点的子节点
if not parent_node:
self.root = left_node
else:
if parent_node.left_child == node:
parent_node.left_child = left_node
else:
parent_node.right_child = left_node
left_node.parent = parent_node
def change_color(self, node):
"""红黑树变色"""
if node.color is 'red':
node.color = 'black'
elif node.color is 'black':
node.color = 'red'
def rb_insert(self, value):
"""红黑树插入"""
node = value if isinstance(value, RBNode) else RBNode(value)
if self.search(self.root, node.data):
return
if self.is_empty():
node.color = 'black'
self.root = node
return
self.insert(self.root, node)
factor_node = node
while True:
parent_node = factor_node.parent
if parent_node.color is 'red':
grandparent_node = parent_node.parent
if parent_node is grandparent_node.left_child:
uncle_node = grandparent_node.right_child
else:
uncle_node = grandparent_node.left_child
# 如果父节点为红节点且叔节点为黑节点
if uncle_node is None or uncle_node and uncle_node.color is 'black':
if parent_node == grandparent_node.left_child:
# 先左旋为左左结果，然后父节点和祖父节点变色，再右旋
if factor_node == parent_node.right_child:
self.left_rotate(parent_node)
self.change_color(factor_node)
else:
self.change_color(parent_node)
self.change_color(grandparent_node)
self.right_rotate(grandparent_node)
elif parent_node == grandparent_node.right_child:
# 先右旋为右右结构，然后父节点和祖父节点变色，再左旋
if factor_node == parent_node.left_child:
self.right_rotate(parent_node)
self.change_color(factor_node)
else:
self.change_color(parent_node)
self.change_color(grandparent_node)
self.left_rotate(grandparent_node)
break
# 如果父节点为红节点且叔节点也为红节点
elif uncle_node and uncle_node.color is 'red':
# 父节点和叔节点变色，祖父节点变色(祖父节点是根节点除外)
self.change_color(parent_node)
self.change_color(uncle_node)
if grandparent_node != self.root:
self.change_color(grandparent_node)
# 祖父节点变成红节点后，将祖父节点作为新的因素节点，判断其父节点，避免不满足特性4
factor_node = grandparent_node
else:
break
def insert(self, root, value):
"""二叉搜索树插入节点-递归"""
node = value if isinstance(value, RBNode) else RBNode(value)
if self.is_empty():
self.root = node
return
if root is None:
root = node
elif node.data < root.data:
root.left_child = self.insert(root.left_child, value)
root.left_child.parent = root
elif node.data > root.data:
root.right_child = self.insert(root.right_child, value)
root.right_child.parent = root
return root
def search(self, root, data):
"""二叉搜索树的查询操作"""
if root is None:
return
if root.data == data:
return root
elif data < root.data:
return self.search(root.left_child, data)
elif data > root.data:
return self.search(root.right_child, data)
def show_tree(self):
if self.is_empty():
print('空二叉树')
return
print('-' * 20)
print("\033[31m{}\033[0m".format(str(self.root.data))) if self.root.color is 'red' else print(str(self.root.data))
self.__print_tree(self.__root)
print('-' * 20)
def __print_tree(self, node, prefix=None):
if prefix is None:
prefix, prefix_left_child = '', ''
else:
prefix = prefix.replace(self.prefix_branch, self.prefix_trunk).replace(self.prefix_leaf, self.prefix_empty)
prefix_left_child = prefix.replace(self.prefix_leaf, self.prefix_empty)
if self.has_child(node):
if node.right_child is not None:
if node.right_child.color is 'red':
print(prefix + self.prefix_branch + self.prefix_right + "\033[31m{}\033[0m".format(str(node.right_child.data)))
else:
print(prefix + self.prefix_branch + self.prefix_right + str(node.right_child.data))
if self.has_child(node.right_child):
self.__print_tree(node.right_child, prefix + self.prefix_branch + ' ')
else:
print(prefix + self.prefix_branch + self.prefix_right)
if node.left_child is not None:
if node.left_child.color is 'red':
print(prefix + self.prefix_leaf + self.prefix_left + "\033[31m{}\033[0m".format(str(node.left_child.data)))
else:
print(prefix + self.prefix_leaf + self.prefix_left + str(node.left_child.data))
if self.has_child(node.left_child):
prefix_left_child += ' '
self.__print_tree(node.left_child, self.prefix_leaf + prefix_left_child)
else:
print(prefix + self.prefix_leaf + self.prefix_left)
def has_child(self, node):
return node.left_child is not None or node.right_child is not None
if __name__ == '__main__':
tree = RBBinaryTree()
data = [50, 77, 55, 29, 10, 30, 66, 18, 80, 51, 90]
for i in data:
# tree.insert(tree.root, i)
tree.rb_insert(i)
tree.show_tree()
# node = tree.search(tree.root, 77)
# tree.left_rotate(node)
# tree.show_tree()
# node = tree.search(tree.root, 80)
# tree.right_rotate(node)
# tree.show_tree()
# tree.change_color(tree.search(tree.root, 50))
# tree.change_color(tree.search(tree.root, 10))
# tree.show_tree()
