红黑树

# 中序遍历
def in_order_tree_walk(node):
    if node:
        in_order_tree_walk(node.leftChild)
        print(node.value, node.color)
        in_order_tree_walk(node.rightChild)


def pre_order_tree_walk(node):
    if node:
        print(node.value, node.color)
        pre_order_tree_walk(node.leftChild)
        pre_order_tree_walk(node.rightChild)


def transplant(tree, node_u, node_v):
    if not node_u.parent:
        tree.root = node_v
    elif node_u == node_u.parent.leftChild:
        node_u.parent.leftChild = node_v
    elif node_u == node_u.parent.rightChild:
        node_u.parent.rightChild = node_v
    if node_v:
        node_v.parent = node_u.parent


def tree_maximum(node):
    temp_node = node
    while temp_node.rightChild:
        temp_node = temp_node.rightChild
    return temp_node


def tree_minimum(node):
    temp_node = node
    while temp_node.leftChild:
        temp_node = temp_node.leftChild
    return temp_node


# 左旋
def left_rotate(tree, x):
    if not x.rightChild:
        return False
    # 假设x的右子结点为y
    y = x.rightChild
    # 将y的左孩子设置为x的右孩子

    x.rightChild = y.leftChild
    # 将y的左孩子的父结点设置为x
    if y.leftChild:
        y.leftChild.parent = x
    y.parent = x.parent
    if not x.parent:
        tree.root = y
    elif x == x.parent.leftChild:
        x.parent.leftChild = y
    else:
        x.parent.rightChild = y
    y.leftChild = x
    x.parent = y


# 右旋
def right_rotate(tree, y):
    if not y.rightChild:
        return False
    x = y.rightChild
    y.rightChild = x.leftChild
    if x.leftChild:
        x.leftChild.parent = y
    x.parent = y.parent
    if not y.parent:
        tree.root = x
    elif y == y.parent.leftChild:
        y.parent.leftChild = x
    else:
        y.parent.rightChild = x
    x.leftChild = y
    y.parent = x


class TreeNode:
    def __init__(self, value):
        self.value = value
        self.parent = None
        self.leftChild = None
        self.rightChild = None
        self.color = 'red'


class RebBlackTree:
    def __init__(self, root=None):
        self.root = root

    # 插入
    def rb_insert(self, z):
        y = None
        x = self.root
        while x:
            y = x
            if z.value == y.value:
                return False
            elif z.value < x.value:
                x = x.leftChild
            else:
                x = x.rightChild
        if not y:
            self.root = z
            z.color = 'black'
        elif z.value < y.value:
            y.leftChild = z
            z.parent = y
        else:
            y.rightChild = z
            z.parent = y
        self.rb_insert_fix_up(z)

    def rb_insert_fix_up(self, z):
        if z.value == self.root.value:
            return
        while z.parent and z.parent.color == 'red':
            if z.parent == z.parent.parent.leftChild:
                y = z.parent.parent.rightChild
                if y and y.color == 'red':
                    z.parent.color = 'black'
                    y.color = 'black'
                    z.parent.parent.color = 'red'
                    z = z.parent.parent
                    continue
                elif z == z.parent.rightChild:
                    left_rotate(self, z.parent)
                    z = z.leftChild
                z.parent.color = 'black'
                z.parent.parent.color = 'red'
                right_rotate(self, z.parent.parent)
                return
            elif z.parent == z.parent.parent.rightChild:
                y = z.parent.parent.leftChild
                if y and y.color == 'red':
                    z.parent.color = 'black'
                    y.color = 'black'
                    z.parent.parent.color = 'red'
                    z = z.parent.parent
                    continue
                elif z == z.parent.leftChild:
                    right_rotate(self, z)
                z.parent.color = 'black'
                z.parent.parent.color = 'red'
                left_rotate(self, z.parent.parent)
                return
        self.root.color = 'black'

    def rb_delete(self, node):
        node_color = node.color
        if not node.leftChild:
            temp_node = node.rightChild
            transplant(self, node, node.rightChild)
        elif not node.rightChild:
            temp_node = node.leftChild
            transplant(self, node, node.leftChild)
        else:
            node_successor = tree_minimum(node.rightChild)
            node_color = node_successor.color
            temp_node = node_successor.rightChild
            if node_successor.parent != node:
                transplant(self, node_successor, node_successor.rightChild)
                node_successor.rightChild = node.rightChild
                node_successor.rightChild.parent = node_successor
            transplant(self, node, node_successor)
            node_successor.leftChild = node.leftChild
            node_successor.leftChild.parent = node_successor
            node_successor.color = node.color
        if node_color == 'black':
            self.delete_fix_up(temp_node)

    def delete_fix_up(self, node):
        while node != self.root and node.color == 'black':
            if node == node.parent.leftChild:
                node_brother = node.parent.rightChild
                if node_brother.color == 'red':
                    node_brother.color = 'black'
                    node.parent.color = 'red'
                    left_rotate(self, node.parent)
                    node_brother = node.parent.rightChild
                if (not node_brother.leftChild or node_brother.leftChild.color == 'black') and (not node_brother.rightChild or node_brother.rightChild.color == 'black'):
                    node_brother.color = 'red'
                    node = node.parent
                else:
                    if not node_brother.rightChild or node_brother.rightChild.color == 'black':
                        node_brother.color = 'red'
                        node_brother.leftChild.color = 'black'
                        right_rotate(self, node_brother)
                        node_brother = node.parent.rightChild
                    node_brother.color = node.parent.color
                    node.parent.color = 'black'
                    node_brother.rightChild.color = 'black'
                    left_rotate(self, node.parent)
                    node = self.root
                # node = self.root
                # break
            else:
                node_brother = node.parent.leftChild
                if node_brother.color == 'red':
                    node_brother.color = 'black'
                    node.parent.color = 'red'
                    left_rotate(self, node.parent)
                    node_brother = node.parent.rightChild
                if (not node_brother.leftChild or node_brother.leftChild.color == 'black') and (not node_brother.rightChild or node_brother.rightChild.color == 'black'):
                    node_brother.color = 'red'
                    node = node.parent
                else:
                    if not node_brother.leftChild or node.leftChild.color == 'black':
                        node.brother.color = 'red'
                        node_brother.rightChild.color = 'black'
                        left_rotate(self, node_brother)
                        node_brother = node.parent.leftChild
                    node_brother.color = node.parent.color
                    node.parent.color = 'black'
                    node_brother.leftChild.color = 'black'
                    right_rotate(self, node.parent)
                    node = self.root
                # node = self.root
                # break
        node.color = 'black'


if __name__ == '__main__':
    Tree = RebBlackTree()
    node_11 = TreeNode(11)
    Tree.rb_insert(node_11)
    node_2 = TreeNode(2)
    Tree.rb_insert(node_2)
    node_14 = TreeNode(14)
    Tree.rb_insert(node_14)
    node_1 = TreeNode(1)
    Tree.rb_insert(node_1)
    node_7 = TreeNode(7)
    Tree.rb_insert(node_7)
    node_15 = TreeNode(15)
    Tree.rb_insert(node_15)
    node_5 = TreeNode(5)
    Tree.rb_insert(node_5)
    node_8 = TreeNode(8)
    Tree.rb_insert(node_8)

    print(Tree.root, Tree.root.value)

    in_order_tree_walk(Tree.root)
    print('+++++++++++++++++++++++++++++++++')
    pre_order_tree_walk(Tree.root)
    print('+++++++++++++++++++++++++++++++++')
    Tree.rb_delete(node_7)

    in_order_tree_walk(Tree.root)
    print('+++++++++++++++++++++++++++++++++')
    pre_order_tree_walk(Tree.root)