最完善红黑树（python）支持可视化(可支持大部分字典操作）

BLACK = 0
RED = 1
# 可视化使用
VC = '│'
HC = '─'
SIZE = 3
RIG = '┌' + HC * SIZE
LEF = '└' + HC * SIZE
SP = chr(32)
IND1 = SP * (SIZE + 1)
IND2 = VC + SP * SIZE


class TreeNode:
    def __init__(self, color=BLACK, key=-1, value=-1, left=None, right=None, p=None):  # 0为黑，1为红
        self.color = color
        self.key = key
        self.left = left
        self.right = right
        self.p = p
        self.value = value

    def __repr__(self):
        return '%s%s%s' % (self.key, '◆' if self.color == BLACK else '◇', self.value)


nil = TreeNode()


class RBTreeNode(TreeNode):
    def __init__(self, key=-1, value=-1):
        super().__init__(RED, key, value, nil, nil, nil)


class REDBLACKTree:
    def __init__(self):
        self.__nil = nil
        self.__root = nil
        self.__l = 0

    def __LEFT_ROTATE(self, x: TreeNode):
        y = x.right
        x.right = y.left
        if y.left != self.__nil:
            y.left.p = x
        y.p = x.p
        if x.p == self.__nil:
            self.__root = y
        elif x == x.p.left:
            x.p.left = y
        else:
            x.p.right = y
        y.left = x
        x.p = y

    def __RIGHT_ROTATE(self, y: TreeNode):
        x = y.left
        y.left = x.right
        if x.right != self.__nil:
            x.right.p = y
        x.p = y.p
        if y.p == self.__nil:
            self.__root = x
        elif y == y.p.left:
            y.p.left = x
        else:
            y.p.right = x
        x.right = y
        y.p = x

    def find_max(self):
        now = self.__root
        while now.right != self.__nil:
            now = now.right
        return now.key

    def find_min(self):
        now = self.__root
        while now.left != self.__nil:
            now = now.left
        return now.key

    def __predecessor(self, x: TreeNode):  # 寻找前继节点
        if x == self.__nil:
            return self.__nil
        if x.left != self.__nil:
            now = x.left
            while now.right != self.__nil:
                now = now.right
            return now
        else:
            p = x.p
            ch = x
            while p != self.__nil and p.left == ch:
                ch = p
                p = p.p
            return p

    def __successor(self, x: TreeNode):
        if x == self.__nil:
            return self.__nil
        if x.right != self.__nil:
            now = x.right
            while now.left != self.__nil:
                now = now.left
            return now
        else:
            p = x.p
            ch = x
            while p != self.__nil and p.right == ch:
                ch = p
                p = p.p
            return p

    def insert(self, key, val):
        if self.__root == self.__nil:
            self.__root = RBTreeNode(key, val)
            self.__root.color = BLACK
            self.__l += 1
            return

        # 寻找插入位置
        now = self.__root
        parent = nil
        while now != self.__nil:
            parent = now
            if now.key == key:  # 元素在树中出现过
                now.value = val
                return
            elif now.key > key:
                now = now.left
            else:
                now = now.right
        self.__l += 1
        # 插入元素
        new = RBTreeNode(key, val)
        new.p = parent
        if key > parent.key:  # 不会出现等于的情况，等于在循环时就退出
            parent.right = new
        else:
            parent.left = new

        # 调整颜色
        self.__adjust_insert(new)

    def __adjust_insert(self, x: TreeNode):
        if x.p.color == BLACK or x == self.__nil or self.__root == x:
            return  # 父亲节点为黑色，不需要调整
        grand_p = x.p.p  # 爷爷节点
        if x.p == grand_p.left:
            if grand_p.right.color == RED:
                grand_p.color = RED
                grand_p.left.color = BLACK
                grand_p.right.color = BLACK
                self.__adjust_insert(grand_p)  # 矛盾上移
            else:
                if x == x.p.right:
                    x.color = BLACK
                    grand_p.color = RED
                    self.__LEFT_ROTATE(x.p)
                else:
                    x.p.color = BLACK
                    grand_p.color = RED
                self.__RIGHT_ROTATE(grand_p)
        else:
            if grand_p.left.color == RED:
                grand_p.color = RED
                grand_p.left.color = BLACK
                grand_p.right.color = BLACK
                self.__adjust_insert(grand_p)  # 矛盾上移
            else:
                if x == x.p.left:
                    x.color = BLACK
                    grand_p.color = RED
                    self.__RIGHT_ROTATE(x.p)
                else:
                    x.p.color = BLACK
                    grand_p.color = RED
                self.__LEFT_ROTATE(grand_p)
        self.__root.color = BLACK  # 根节点一定为黑

    def remove(self, key):
        now = self.__root
        while now != self.__nil:
            if now.key == key:  # 元素在树中出现过,返回元素值
                break
            elif now.key > key:
                now = now.left
            else:
                now = now.right
        if now == self.__nil:
            return  # 不存在该节点
        else:
            self.__l -= 1
            value = now.value
            # 删除节点
            self.__deleteNode(now)
            return value

    def __deleteNode(self, x: TreeNode):
        if x.left != self.__nil and x.right != self.__nil:
            __successor = self.__successor(x)
            x.key = __successor.key
            x.value = __successor.value
            self.__deleteNode(__successor)
        else:
            replacement = self.__nil
            if x.left != self.__nil:
                replacement = x.left
            elif x.right != self.__nil:
                replacement = x.right

            if replacement == self.__nil:  # 节点为叶子节点
                if x.p == self.__nil:  # 待删除节点为根节点
                    self.__root = self.__nil
                else:
                    # 先进行调整
                    if x.color == BLACK:
                        # 需要进行调整
                        self.__adjust_remove(x)
                    if x == x.p.left:
                        x.p.left = self.__nil
                    else:
                        x.p.right = self.__nil
                    x.p = self.__nil
            else:
                replacement.p = x.p
                if x.p == self.__nil:
                    self.__root = replacement
                else:
                    if x == x.p.left:
                        x.p.left = replacement
                    else:
                        x.p.right = replacement
                if x.color == BLACK:
                    # 需要进行调整
                    self.__adjust_remove(replacement)
            # 将x彻底断开
            x.right = self.__nil
            x.left = self.__nil
            x.p = self.__nil

    def __adjust_remove(self, x: TreeNode):
        if x != self.__root and x.color == BLACK:
            if x == x.p.left:
                brother = x.p.right
                if brother.color == RED:
                    brother.color = BLACK
                    x.p.color = RED
                    self.__LEFT_ROTATE(x.p)
                    brother = x.p.right
                # 情况二，找兄弟接，兄弟没得借
                if brother.left == self.__nil and brother.right == self.__nil:
                    brother.color = RED
                    self.__adjust_remove(x.p)
                else:
                    # 分两种情况，兄弟节点为3或4节点
                    if brother.right == self.__nil:
                        brother.left.color = BLACK
                        brother.color = RED
                        self.__RIGHT_ROTATE(brother)
                        brother = x.p.right
                    brother.color = x.p.color
                    x.p.color = BLACK
                    brother.right.color = BLACK
                    self.__LEFT_ROTATE(x.p)
            else:
                brother = x.p.left
                if brother.color == RED:
                    brother.color = BLACK
                    x.p.color = RED
                    self.__RIGHT_ROTATE(x.p)
                    brother = x.p.left
                # 情况二，找兄弟接，兄弟没得借
                if brother.right == self.__nil and brother.left == self.__nil:
                    brother.color = RED
                    self.__adjust_remove(x.p)
                else:
                    # 分两种情况，兄弟节点为3或4节点
                    if brother.left == self.__nil:  #
                        brother.right.color = BLACK
                        brother.color = RED
                        self.__LEFT_ROTATE(brother)
                        brother = x.p.left
                    brother.color = x.p.color
                    x.p.color = BLACK
                    brother.left.color = BLACK
                    self.__RIGHT_ROTATE(x.p)

        else:
            x.color = BLACK  # 情况一，替代节点为红色，直接设为黑色

    def keys(self):
        if self.__root == self.__nil:
            return
        WHITE, GRAY = 0, 1
        stack = [(WHITE, self.__root)]
        while stack:
            color, node = stack.pop()
            if node == self.__nil:
                continue
            if color == WHITE:
                stack.append((WHITE, node.right))
                stack.append((GRAY, node))
                stack.append((WHITE, node.left))
            else:
                yield node.key

    def values(self):
        if self.__root == self.__nil:
            return
        WHITE, GRAY = 0, 1
        stack = [(WHITE, self.__root)]
        while stack:
            color, node = stack.pop()
            if node == self.__nil:
                continue
            if color == WHITE:
                stack.append((WHITE, node.right))
                stack.append((GRAY, node))
                stack.append((WHITE, node.left))
            else:
                yield node.value

    def items(self):
        if self.__root == self.__nil:
            return
        WHITE, GRAY = 0, 1
        stack = [(WHITE, self.__root)]
        while stack:
            color, node = stack.pop()
            if node == self.__nil:
                continue
            if color == WHITE:
                stack.append((WHITE, node.right))
                stack.append((GRAY, node))
                stack.append((WHITE, node.left))
            else:
                yield node.key, node.value

    def empty(self):
        return self.__root == self.__nil

    def __len__(self):
        return self.__l

    def __getitem__(self, key):
        now = self.__root
        while now != self.__nil:
            if now.key == key:  # 元素在树中出现过,返回元素值
                return now.value
            elif now.key > key:
                now = now.left
            else:
                now = now.right
        return 0

    def __delitem__(self, key):
        self.remove(key)

    def __setitem__(self, key, value):
        self.insert(key, value)

    def __repr__(self):
        return '\n'.join(self.__graph())

    def __iter__(self):
        if self.__root == self.__nil:
            return
        WHITE, GRAY = 0, 1
        stack = [(WHITE, self.__root)]
        while stack:
            color, node = stack.pop()
            if node == self.__nil:
                continue
            if color == WHITE:
                stack.append((WHITE, node.right))
                stack.append((GRAY, node))
                stack.append((WHITE, node.left))
            else:
                yield node.key

    def __contains__(self, key):
        now = self.__root
        while now != self.__nil:
            if now.key == key:  # 元素在树中出现过
                return True
            elif now.key > key:
                now = now.left
            else:
                now = now.right
        return False

    def __graph(self, x=False, prefix=''):
        """beautifully print RBtree, big key node first"""
        if x is False:
            x = self.__root
        if x is not self.__nil:
            p = x.p
            last_prefix = ''
            if p is not self.__nil:
                pp = p.p
                last_prefix = LEF if p.left is x else RIG
                if pp is not self.__nil:
                    if (pp.left is p) is (p.left is x):
                        prefix = prefix + IND1
                    else:
                        prefix = prefix + IND2
            yield from self.__graph(x.right, prefix)
            yield '%s%s%s' % (prefix, last_prefix, x)
            yield from self.__graph(x.left, prefix)