数据结构 - 红黑树

#!/usr/bin/env python3.3
# -*- coding:utf-8 -*-
# Copyright 2013'''
满足如下红黑性质的二叉查找树：
1）节点是红色或黑色
2）根是黑色
3）所有叶子都是黑色（叶子是NIL节点）
4）每个红色节点的两个子节点都是黑色（从每个叶子到根的所有路径上不能有两个连续的红色节点）
5）从任一节点到其每个叶子的所有简单路径都包含相同数目的黑色节点
'''class Node:def __init__(self, key):self.key = keyself.left = Noneself.right = Noneself.color = Noneself.parent = Nonedef grand_parent(self):if self.parent is not None:return self.parent.parentdef uncle(self):if self.grand_parent() is not None:if self.parent == self.grand_parent().left:return self.grand_parent().rightelse:return self.grand_parent().leftdef sibling(self):if self.parent is not None:if self.parent.left == self:return self.parent.rightelse:return self.parent.leftdef __str__(self):return str(self.key)class RBTree:def __init__(self):self.root = Nonedef serach(self, key):node = self.rootwhile node is not None:if key == node.key:return nodeelif key < node.key:node = node.leftelse:node = node.rightdef insert(self, key):# 首先标记为红色node = Node(key)node.color = 'RED'if self.root is None:self.root = nodeelse:root = self.rootwhile root is not None:if key < root.key:if root.left is None:root.left = nodenode.parent = rootbreakelse:root = root.left                else:if root.right is None:root.right = nodenode.parent = rootbreakelse:root = root.rightself._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.rightwhile successor.left is not None:successor = successor.left node.key = successor.keynode = successorif node.left is not None:# 要删除的节点没有右儿子，左儿子（红色）node.left.color = 'BLACK'node.left.parent = node.parentif node.parent is not None:if node.parent.left == node:node.parent.left = node.leftelse:node.parent.right = node.leftelse:self.root = node.leftelif node.right is not None:# 要删除的节点没有左儿子，右儿子（红色）node.right.color = 'BLACK'node.right.parent = node.parentif node.parent is not None:if node.parent.left == node:node.parent.left = node.rightelse:node.parent.right = node.rightelse:self.root = node.rightelse:# 要删除的节点为叶子节点if node.color == 'BLACK':self._delete_fixup(node)if node.parent is not None:if node.parent.left == node:node.parent.left = Noneelse:node.parent.right = Noneelse:self.root = Nonedef _insert_fixup(self, node):if node.parent is None:# 1位于树根上node.color = 'BLACK'elif node.parent.color == 'BLACK':# 2父节点P（黑色）passelif 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.leftnode.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.rightnode.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.colornode.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.colornode.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.rightright_left = right.left# 处理右子的左子节点if right_left is not None:right_left.parent = node# 处理右子节点right.left = noderight.parent = node.parent# 处理父节点if node.parent is not None:if node.parent.left == node:node.parent.left = rightelse:node.parent.right = rightelse:self.root = right# 处理当前节点node.parent = rightnode.right = right_leftdef _rotate_right(self, node):left = node.leftleft_right = left.right# 处理左子的右子节点if left_right is not None:left_right.parent = node# 处理左子节点left.right = nodeleft.parent = node.parent# 处理父节点if node.parent is not None:if node.parent.right == node:node.parent.right = leftelse:node.parent.left = leftelse:self.root = left# 处理当前节点node.parent = leftnode.left = left_rightdef 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')breaknode_list.append(node.left)if node.right is not None:if node.key > node.right.key:print('sort_error')breaknode_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')breaknode_list.append(node.left)if node.right is not None:if node.color == 'RED' and node.right.color == 'RED':print('double_red')breaknode_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 = 0while len(leaf_list) > 0:node = leaf_list.pop(0)node_depth = 0while node is not None:if node.color == 'BLACK':node_depth += 1node = node.parentif black_depth == 0:black_depth = node_depthelse:if node_depth != black_depth:print('black_depth_error')breakprint('ok')else:print('empty')