RBTree

rbtree.py

from enum import Enum
import random


class Color(Enum):
    RED = "red"
    BLACK = "black"


class RBNode:
    data: int
    color: Color
    left: RBNode
    right: RBNode
    parent: RBNode
    count: int = 1

    def __init__(self, data):
        self.data = data
        self.color = Color.RED  # New nodes are red by default
        self.left = None
        self.right = None
        self.parent = None

    def __repr__(self):
        return f"RBNode(data={self.data}, color={self.color}, left={self.left.data if self.left else None}, right={self.right.data if self.right else None}, count={self.count})"


class RBTree:
    def __init__(self):
        self.NIL = RBNode(0)  # Sentinel node
        self.NIL.color = Color.BLACK
        self.root = self.NIL

    def left_rotate(self, node: RBNode):
        y = node.right

        if node.parent is None:
            self.root = y
        elif node == node.parent.left:
            node.parent.left = y
        else:
            node.parent.right = y

        y.parent = node.parent

        node.right = y.left
        if y.left != self.NIL:
            y.left.parent = node

        y.left = node
        node.parent = y
        pass

    def right_rotate(self, node: RBNode):
        y = node.left

        if node.parent is None:
            self.root = y
        elif node == node.parent.right:
            node.parent.right = y
        else:
            node.parent.left = y

        y.parent = node.parent

        node.left = y.right
        if y.right != self.NIL:
            y.right.parent = node

        y.right = node
        node.parent = y
        pass

    def insert(self, data):
        new_node = RBNode(data)
        new_node.left = self.NIL
        new_node.right = self.NIL

        parent = None
        current = self.root

        while current != self.NIL:
            parent = current
            if new_node.data == current.data:
                current.count += 1
                return current
            elif new_node.data < current.data:
                current = current.left
            else:
                current = current.right

        new_node.parent = parent

        # insert the new node
        if parent is None:  # Tree was empty
            self.root = new_node
        elif new_node.data < parent.data:
            parent.left = new_node
        else:
            parent.right = new_node

        self._fix_insert(new_node)
        return new_node

    def _fix_insert(self, new_node: RBNode):
        r"""
        Fix the tree after insertion to maintain Red-Black properties.

                       G
                      / \
                     P   U
                    /\  / \
                   N  N N  N
        """

        while new_node.parent and new_node.parent.color == Color.RED:
            node_p = new_node.parent
            node_g = node_p.parent if node_p else None
            node_u = (
                (node_g.right if node_p == node_g.left else node_g.left)
                if node_g
                else None
            )
            if node_p == node_g.left:
                if node_u and node_u.color == Color.RED:
                    # Case 1: Uncle is red
                    node_p.color = Color.BLACK
                    node_u.color = Color.BLACK
                    node_g.color = Color.RED
                    # node_g 被染红，作为 new_node 向上考虑
                    new_node = node_g
                else:
                    # Case 2: Uncle is black and new_node is right child
                    if new_node == node_p.right:
                        new_node = node_p
                        self.left_rotate(new_node)
                        node_p = new_node.parent
                    # Case 3: Uncle is black and new_node is left child
                    node_p.color = Color.BLACK
                    node_g.color = Color.RED
                    self.right_rotate(node_g)
            else:
                if node_u and node_u.color == Color.RED:
                    # Case 1: Uncle is red
                    node_p.color = Color.BLACK
                    node_u.color = Color.BLACK
                    node_g.color = Color.RED
                    # node_g 被染红，作为 new_node 向上考虑
                    new_node = node_g
                else:
                    # Case 2: Uncle is black and new_node is left child
                    if new_node == node_p.left:
                        new_node = node_p
                        self.right_rotate(new_node)
                        node_p = new_node.parent
                    # Case 3: Uncle is black and new_node is right child
                    node_p.color = Color.BLACK
                    node_g.color = Color.RED
                    self.left_rotate(node_g)

        self.root.color = Color.BLACK

    def _transplant(self, u: RBNode, v: RBNode):
        """用树 v 替换树 u

        注意不是交换，直接把 v 接在 u 的父节点上
        """
        if u.parent is None:
            self.root = v
        elif u == u.parent.left:  # u 是左子树
            u.parent.left = v
        else:  # u 是右子树
            u.parent.right = v
        v.parent = u.parent

    def _delete_fix(self, node: RBNode):
        while node != self.root and node.color == Color.BLACK:
            if node == node.parent.left:
                sibling = node.parent.right
                # Case 1: sibling is red
                if sibling.color == Color.RED:
                    sibling.color = Color.BLACK
                    node.parent.color = Color.RED
                    self.left_rotate(node.parent)
                    sibling = node.parent.right

                # Case 2: sibling is black and both of sibling's children are black
                if (
                    sibling.left.color == Color.BLACK
                    and sibling.right.color == Color.BLACK
                ):
                    sibling.color = Color.RED
                    # node upward to parent, continue checking
                    node = node.parent
                else:
                    # Case 3: sibling is black and sibling's right child is black
                    if sibling.right.color == Color.BLACK:
                        sibling.left.color = Color.BLACK
                        sibling.color = Color.RED
                        self.right_rotate(sibling)
                        sibling = node.parent.right

                    # Case 4: sibling is black and sibling's right child is red
                    sibling.color = node.parent.color
                    node.parent.color = Color.BLACK
                    sibling.right.color = Color.BLACK
                    self.left_rotate(node.parent)
                    break
            else:
                sibling = node.parent.left
                if sibling.color == Color.RED:
                    sibling.color = Color.BLACK
                    node.parent.color = Color.RED
                    self.right_rotate(node.parent)
                    sibling = node.parent.left

                if (
                    sibling.left.color == Color.BLACK
                    and sibling.right.color == Color.BLACK
                ):
                    sibling.color = Color.RED
                    node = node.parent
                else:
                    if sibling.left.color == Color.BLACK:
                        sibling.right.color = Color.BLACK
                        sibling.color = Color.RED
                        self.left_rotate(sibling)
                        sibling = node.parent.left

                    sibling.color = node.parent.color
                    node.parent.color = Color.BLACK
                    sibling.left.color = Color.BLACK
                    self.right_rotate(node.parent)
                    break

        node.color = Color.BLACK

    def _delete_node(self, z: RBNode):
        y = z
        y_original_color = y.color

        # 情况 1 & 2: 只有一个孩子或没有孩子
        if z.left == self.NIL:
            r"""
                    z            z
                   / \    or    / \
                  N   X        N   N
            """
            x = z.right
            self._transplant(z, z.right)
            # 最后根据 y 的颜色是否为黑色 _delete_fix(x)
        elif z.right == self.NIL:
            r"""
                    z
                   / \
                  X   N
            """
            x = z.left
            self._transplant(z, z.left)
            # 最后根据 y 的颜色是否为黑色 _delete_fix(x)
        else:
            # 情况 3: 有两个孩子，找后继节点 (Successor)
            r"""
                  z
                 /  \
               z.L  ...
                     \
                      y
                     / \ 
                    N   x
            """
            y = z.right
            while y.left != self.NIL:
                y = y.left

            y_original_color = y.color
            x = y.right  # 后继节点 y 最多只有一个右孩子
            # 最后从最底下开始 _delete_fix(x)

            if y.parent == z:
                r"""
                    z
                   / \
                  z.L y
                     / \ 
                    N   x
                """
                x.parent = y  # 如果 y 是 z 的直接孩子，建立 parent 连接
            else:
                # 将 y 的右子树替换 y 的位置
                self._transplant(y, y.right)
                y.right = z.right
                y.right.parent = y

            # 将 z 替换为 y
            self._transplant(z, y)
            y.left = z.left
            y.left.parent = y

            y.color = z.color  # y 继承 z 的颜色，保持结构稳定

        # 如果被删除/移动的是黑色节点，必须修复
        if y_original_color == Color.BLACK:
            self._delete_fix(x)  # 修复的节点 x 是被剩下的的唯一子节点

    def delete(self, data: int):
        # 寻找节点的逻辑保持不变
        current = self.root
        while current != self.NIL:
            if data == current.data:
                if current.count > 1:
                    current.count -= 1
                    return True
                else:
                    self._delete_node(current)
                    return True
            elif data < current.data:
                current = current.left
            else:
                current = current.right
        return False


def verify_rbtree(tree: RBTree):
    if tree.root == tree.NIL:
        return True  # 空树也是红黑树

    # 规则 2: 根节点必须是黑色 (Root Property)
    if tree.root.color != Color.BLACK:
        raise Exception("Root must be black")

    def helper(node: RBNode) -> int:
        # 规则 3: NIL 节点是黑色 (NIL Property)
        if node == tree.NIL:
            return 1  # 贡献 1 个黑高度

        # 规则 4: 红色节点不能有红色子节点 (Red Property)
        if node.color == Color.RED:
            if node.left.color == Color.RED or node.right.color == Color.RED:
                raise Exception(f"Red node {node.data} has red children")

        # 二叉搜索树性质验证
        if node.left != tree.NIL and node.left.data > node.data:
            raise Exception(f"BST violation: {node.left.data} > {node.data}")
        if node.right != tree.NIL and node.right.data < node.data:
            raise Exception(f"BST violation: {node.right.data} < {node.data}")

        # 规则 5: 递归检查左右子树的黑高 (Black Height Property)
        left_black_height = helper(node.left)
        right_black_height = helper(node.right)

        if left_black_height != right_black_height:
            raise Exception(
                f"Black height mismatch at node {node.data}: "
                f"left {left_black_height}, right {right_black_height}"
            )

        # 返回当前节点的黑高贡献
        return left_black_height + (1 if node.color == Color.BLACK else 0)

    # 开始执行递归检查
    helper(tree.root)
    return True


if __name__ == "__main__":
    tree = RBTree()
    values = [value for value in range(100)]
    random.shuffle(values)
    for value in values:
        tree.insert(value)
        assert verify_rbtree(tree) == True
    for value in values:
        tree.delete(value)
        assert verify_rbtree(tree) == True
    print("All tests passed!")