紅黑樹Python實現

# 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()

教材: 算法導論(第三版)