红黑树的Python实现示例

红黑树是一种自平衡二叉查找树，具有多种特性确保了它在最坏情况下的操作时间复杂度为O(log n)。本文将通过Python代码实现一个简单的红黑树，并展示其基本的操作方法。

1. 红黑树的基本性质

红黑树有以下5个重要性质：

每一个节点或者是红色，或者是黑色。
根节点是黑色的。
所有叶子节点（NIL或空节点）都是黑色的。
如果一个节点是红色的，则它的两个子节点都是黑色的。（从每个叶到根的所有路径上不能有两个连续的红色节点）
从任一节点到其每个叶子的所有简单路径都包含相同数目的黑色节点。

2. 红黑树类定义

首先定义红黑树的基本结构，并初始化一些属性，如颜色、父节点和子节点等。

class Node:
    def __init__(self, key):
        self.key = key
        self.color = 'RED'  # 初始颜色为红色
        self.parent = None
        self.left = None
        self.right = None

class RedBlackTree:
    def __init__(self):
        self.NIL = Node(None)  # 空节点
        self.NIL.color = 'BLACK'
        self.root = self.NIL

3. 插入操作

插入新元素时，我们首先将它作为红色节点插入到二叉树中。然后通过一系列旋转和着色调整来恢复红黑树的性质。

def insert(self, key):
    node = Node(key)
    node.left = self.NIL
    node.right = self.NIL

    parent = None
    current = self.root
    while current != self.NIL:
        parent = current
        if node.key < current.key:
            current = current.left
        else:
            current = current.right

    node.parent = parent
    if parent is None:
        self.root = node  # 树为空，直接插入为根节点
    elif node.key < parent.key:
        parent.left = node
    else:
        parent.right = node
    node.color = 'RED'  # 插入时颜色设为红色

    self.fix_insert(node)

4. 修复插入操作

需要对插入后的红黑树进行调整，使其满足红黑树的性质。

def fix_insert(self, k):
    while k.parent.color == 'RED':
        if k.parent == k.parent.parent.left:
            u = k.parent.parent.right  # 叔叔节点
            if u.color == 'RED':  # 情况1：叔叔是红色的
                u.color = 'BLACK'
                k.parent.color = 'BLACK'
                k.parent.parent.color = 'RED'
                k = k.parent.parent
            else:
                if k == k.parent.right:  # 情况2：k是右孩子，且父节点是左子树
                    k = k.parent
                    self.left_rotate(k)
                # 情况3：k是左孩子，调整后进行一次右旋
                k.parent.color = 'BLACK'
                k.parent.parent.color = 'RED'
                self.right_rotate(k.parent.parent)
        else:  # 类似的情况2和情况3（叔叔在右子树中）
            u = k.parent.parent.left

            if u.color == 'RED':
                u.color = 'BLACK'
                k.parent.color = 'BLACK'
                k.parent.parent.color = 'RED'
                k = k.parent.parent
            else:
                if k == k.parent.left:  # 情况2，k是左孩子
                    k = k.parent
                    self.right_rotate(k)
                # 情况3：k是右孩子
                k.parent.color = 'BLACK'
                k.parent.parent.color = 'RED'
                self.left_rotate(k.parent.parent)

        if k == self.root:
            break

    self.root.color = 'BLACK'  # 根节点总是黑色的

5. 左旋操作

左旋是红黑树中的一种基本旋转方法，用于调整红黑树中的节点位置。

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 is None:  # 当x是根节点时，调整根节点
        self.root = y
    elif x == x.parent.left:
        x.parent.left = y
    else:
        x.parent.right = y
    y.left = x
    x.parent = y

6. 右旋操作

右旋与左旋类似，用于调整红黑树中的节点位置。

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

    y.parent = x.parent
    if x.parent is None:  # 当x是根节点时，调整根节点
        self.root = y
    elif x == x.parent.right:
        x.parent.right = y
    else:
        x.parent.left = y
    y.right = x
    x.parent = y

总结

通过上述实现，我们已经完成了红黑树的基本插入操作以及相关的修复逻辑。这些基本操作保证了插入元素后的红黑树仍然满足性质要求。红黑树的实现不仅展示了平衡二叉查找树的思想，也体现了对算法细节的关注和处理。