# 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()
教材: 算法導論(第三版)