红黑树的具体代码实现

52 阅读1分钟
class Node:
    def __init__(self, key, left=None, right=None, parent=None, color='red'):
        self.key = key
        self.left = left
        self.right = right
        self.parent = parent
        self.color = color

    def __repr__(self):
        return f"Node({self.key}, {self.color})"

class RedBlackTree:
    def __init__(self):
        self.NIL = Node(None, color='black')  # 哨兵节点,所有空指针都指向它
        self.root = self.NIL

    def left_rotate(self, x):
        y = x.right
        x.right = y.left
        if y.left != self.NIL:
            y.left.parent = x
        y.parent = x.parent
        if x.parent == self.NIL:
            self.root = y
        elif x == x.parent.left:
            x.parent.left = y
        else:
            x.parent.right = y
        y.left = x
        x.parent = y

    def right_rotate(self, y):
        x = y.left
        y.left = x.right
        if x.right != self.NIL:
            x.right.parent = y
        x.parent = y.parent
        if y.parent == self.NIL:
            self.root = x
        elif y == y.parent.right:
            y.parent.right = x
        else:
            y.parent.left = x
        x.right = y
        y.parent = x

    def insert_fixup(self, z):
        while z.parent.color == 'red':
            if z.parent == z.parent.parent.left:
                y = z.parent.parent.right
                if y.color == 'red':
                    z.parent.color = 'black'
                    y.color = 'black'
                    z.parent.parent.color = 'red'
                    z = z.parent.parent
                else:
                    if z == z.parent.right:
                        z = z.parent
                        self.left_rotate(z)
                    z.parent.color = 'black'
                    z.parent.parent.color = 'red'
                    self.right_rotate(z.parent.parent)
            else:
                # 对称的情况
                y = z.parent.parent.left
                if y.color == 'red':
                    z.parent.color = 'black'
                    y.color = 'black'
                    z.parent.parent.color = 'red'
                    z = z.parent.parent
                else:
                    if z == z.parent.left:
                        z = z.parent
                        self.right_rotate(z)
                    z.parent.color = 'black'
                    z.parent.parent.color = 'red'
                    self.left_rotate(z.parent.parent)

    def insert(self, key):
        z = Node(key)
        y = self.NIL
        x = self.root

        while x != self.NIL:
            y = x
            if z.key < x.key:
                x = x.left
            else:
                x = x.right

        z.parent = y

        if y == self.NIL:
            self.root = z
        elif z.key < y.key:
            y.left = z
        else:
            y.right = z

        z.left = self.NIL
        z.right = self.NIL
        z.color = 'red'

        self.insert_fixup(z)

# 使用示例
rbt = RedBlackTree()
rbt.insert(7)
rbt.insert(3)
rbt.insert(18)
rbt.insert(10)
rbt.insert(22)
rbt.insert(8)
rbt.insert(11)
rbt.insert(26)

# 注意:这里没有打印树的函数,因为直接打印树结构比较复杂。
# 你需要实现一个遍历函数(如中序遍历)来查看树的内容。

这个实现包括了插入操作和插入后的修复操作(insert_fixup),以及左旋和右旋操作。注意,这个实现假设了所有空指针都指向一个哨兵节点NIL,它的颜色总是黑色,这有助于简化边界条件的处理。

要查看树的内容,你需要实现一个遍历函数(如中序遍历),因为红黑树是一个有序的二叉树,中序遍历将按升序输出所有键。