#!/usr/bin/env python3.3
# -*- coding:utf-8 -*-
# Copyright 2013
'''
满足如下红黑性质的二叉查找树：
1）节点是红色或黑色
2）根是黑色
3）所有叶子都是黑色（叶子是NIL节点）
4）每个红色节点的两个子节点都是黑色（从每个叶子到根的所有路径上不能有两个连续的红色节点）
5）从任一节点到其每个叶子的所有简单路径都包含相同数目的黑色节点
'''
class Node:
def __init__(self, key):
self.key = key
self.left = None
self.right = None
self.color = None
self.parent = None
def grand_parent(self):
if self.parent is not None:
return self.parent.parent
def uncle(self):
if self.grand_parent() is not None:
if self.parent == self.grand_parent().left:
return self.grand_parent().right
else:
return self.grand_parent().left
def sibling(self):
if self.parent is not None:
if self.parent.left == self:
return self.parent.right
else:
return self.parent.left
def __str__(self):
return str(self.key)
class RBTree:
def __init__(self):
self.root = None
def serach(self, key):
node = self.root
while node is not None:
if key == node.key:
return node
elif key < node.key:
node = node.left
else:
node = node.right
def insert(self, key):
# 首先标记为红色
node = Node(key)
node.color = 'RED'
if self.root is None:
self.root = node
else:
root = self.root
while root is not None:
if key < root.key:
if root.left is None:
root.left = node
node.parent = root
break
else:
root = root.left
else:
if root.right is None:
root.right = node
node.parent = root
break
else:
root = root.right
self._insert_fixup(node)
def delete(self, key):
node = self.search(key)
if node is not None:
if node.left is not None and node.right is not None:
# 把要删除的节点替换为其后继节点
successor = node.right
while successor.left is not None:
successor = successor.left
node.key = successor.key
node = successor
if node.left is not None:
# 要删除的节点没有右儿子，左儿子（红色）
node.left.color = 'BLACK'
node.left.parent = node.parent
if node.parent is not None:
if node.parent.left == node:
node.parent.left = node.left
else:
node.parent.right = node.left
else:
self.root = node.left
elif node.right is not None:
# 要删除的节点没有左儿子，右儿子（红色）
node.right.color = 'BLACK'
node.right.parent = node.parent
if node.parent is not None:
if node.parent.left == node:
node.parent.left = node.right
else:
node.parent.right = node.right
else:
self.root = node.right
else:
# 要删除的节点为叶子节点
if node.color == 'BLACK':
self._delete_fixup(node)
if node.parent is not None:
if node.parent.left == node:
node.parent.left = None
else:
node.parent.right = None
else:
self.root = None
def _insert_fixup(self, node):
if node.parent is None:
# 1位于树根上
node.color = 'BLACK'
elif node.parent.color == 'BLACK':
# 2父节点P（黑色）
pass
elif node.uncle() is not None and node.uncle().color == 'RED':
# 3父节点P（红色），叔父节点U（红色）
node.parent.color = 'BLACK'
node.uncle().color = 'BLACK'
node.grand_parent().color = 'RED'
self._insert_fixup(node.grand_parent())
else:
# 4父节点P（红色），叔父节点U（缺少或黑色）
if node.grand_parent().left == node.parent:
# 父节点是祖父节点的左儿子
if node.parent.right == node:
# 当前节点是父节点的右儿子
self._rotate_left(node.parent)
node = node.left
node.parent.color = 'BLACK'
node.grand_parent().color = 'RED'
self._rotate_right(node.grand_parent())
else:
# 父节点是祖父节点的右儿子
if node.parent.left == node:
# 当前节点是父节点的左儿子
self._rotate_right(node.parent)
node = node.right
node.parent.color = 'BLACK'
node.grand_parent().color = 'RED'
self._rotate_left(node.grand_parent())
def _delete_fixup(self, node):
if node.color == 'RED' or node.parent is None:
node.color = 'BLACK'
else:
if node.sibling().color == 'RED':
# 兄弟节点（红色），旋转使兄弟节点变为（黑色）
node.sibling().color = 'BLACK'
node.parent.color = 'RED'
if node.parent.left == node:
self._rotate_left(node.parent)
else:
self._rotate_right(node.parent)
if (node.sibling().left is None or node.sibling().left.color == 'BLACK') and (
node.sibling().right is None or node.sibling().right.color == 'BLACK'):
# 兄弟节点（黑色），两个侄子节点（缺少或黑色）
node.sibling().color = 'RED'
# Token向上传递（该Token下的所有子孙，黑高度缺少一）
self._delete_fixup(node.parent)
else:
# 兄弟节点（黑色），至少一个侄子节点（红色）
if node.parent.left == node:
if node.sibling().right is None or node.sibling().right.color == 'BLACK':
# 左侄子节点（红色），右侄子节点（缺少或黑色）
node.sibling().left.color = 'BLACK'
node.sibling().color = 'RED'
self._rotate_right(node.sibling())
# 左侄子节点（缺少或黑色），右侄子节点（红色）
node.sibling().color = node.parent.color
node.parent.color = 'BLACK'
if node.sibling().right is not None:
node.sibling().right.color = 'BLACK'
self._rotate_left(node.parent)
else:
if node.sibling().left is None or node.sibling().left.color == 'BLACK':
# 左侄子节点（缺少或黑色），右侄子节点（红色）
node.sibling().right.color = 'BLACK'
node.sibling().color = 'RED'
self._rotate_left(node.sibling())
# 左侄子节点（红色），右侄子节点（缺少或黑色）
node.sibling().color = node.parent.color
node.parent.color = 'BLACK'
if node.sibling().left is not None:
node.sibling().left.color = 'BLACK'
self._rotate_right(node.parent)
def _rotate_left(self, node):
right = node.right
right_left = right.left
# 处理右子的左子节点
if right_left is not None:
right_left.parent = node
# 处理右子节点
right.left = node
right.parent = node.parent
# 处理父节点
if node.parent is not None:
if node.parent.left == node:
node.parent.left = right
else:
node.parent.right = right
else:
self.root = right
# 处理当前节点
node.parent = right
node.right = right_left
def _rotate_right(self, node):
left = node.left
left_right = left.right
# 处理左子的右子节点
if left_right is not None:
left_right.parent = node
# 处理左子节点
left.right = node
left.parent = node.parent
# 处理父节点
if node.parent is not None:
if node.parent.right == node:
node.parent.right = left
else:
node.parent.left = left
else:
self.root = left
# 处理当前节点
node.parent = left
node.left = left_right
def check_rbtree(rbt):
if rbt.root is not None:
# 是一棵二叉查找树
node_list = [rbt.root]
while len(node_list) > 0:
node = node_list.pop(0)
if node.left is not None:
if node.key < node.left.key:
print('sort_error')
break
node_list.append(node.left)
if node.right is not None:
if node.key > node.right.key:
print('sort_error')
break
node_list.append(node.right)
# 根是黑色
if rbt.root.color != 'BLACK':
print('root_color_error')
# 每个红色节点的两个子节点都是黑色
node_list = [rbt.root]
while len(node_list) > 0:
node = node_list.pop(0)
if node.left is not None:
if node.color == 'RED' and node.left.color == 'RED':
print('double_red')
break
node_list.append(node.left)
if node.right is not None:
if node.color == 'RED' and node.right.color == 'RED':
print('double_red')
break
node_list.append(node.right)
# 黑高度相等
node_list = [rbt.root]
leaf_list = []
while len(node_list) > 0:
node = node_list.pop(0)
if node.left is not None:
node_list.append(node.left)
if node.right is not None:
node_list.append(node.right)
if node.left is None or node.right is None:
leaf_list.append(node)
black_depth = 0
while len(leaf_list) > 0:
node = leaf_list.pop(0)
node_depth = 0
while node is not None:
if node.color == 'BLACK':
node_depth += 1
node = node.parent
if black_depth == 0:
black_depth = node_depth
else:
if node_depth != black_depth:
print('black_depth_error')
break
print('ok')
else:
print('empty')
