学习红黑树的概念及利用python代码实现红黑树三种操作和扩张

红黑树是一种自平衡二叉查找树，它保证了在最坏情况下基本动态集合操作的时间复杂度为O(log n)。红黑树的每个节点都有一个颜色属性，可以是红色或黑色。红黑树满足以下5个性质：
1. 每个节点要么是红色，要么是黑色。
2. 根节点是黑色。
3. 每个叶子节点（NIL节点，空节点）是黑色的。
4. 如果一个节点是红色的，则它的两个子节点都是黑色的。
5. 对于每个节点，从该节点到其所有后代叶子节点的简单路径上，均包含相同数目的黑色节点。

红黑树的三种操作包括插入、删除和查找。插入和删除操作需要通过旋转和颜色修改来保证红黑树的性质不被破坏。以下是Python代码实现红黑树的三种操作和扩张：

class Node:
    def __init__(self, key):
        self.key = key
        self.left = None
        self.right = None
        self.parent = None
        self.color = 1

class RedBlackTree:
    def __init__(self):
        self.null_node = Node(0)
        self.null_node.color = 0
        self.null_node.left = None
        self.null_node.right = None
        self.root = self.null_node

    def left_rotate(self, x):
        y = x.right
        x.right = y.left
        if y.left != self.null_node:
            y.left.parent = x
        y.parent = x.parent
        if x.parent == self.null_node:
            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, x):
        y = x.left
        x.left = y.right
        if y.right != self.null_node:
            y.right.parent = x
        y.parent = x.parent
        if x.parent == self.null_node:
            self.root = y
        elif x == x.parent.right:
            x.parent.right = y
        else:
            x.parent.left = y
        y.right = x
        x.parent = y

    def insert(self, key):
        node = Node(key)
        node.parent = None
        node.key = key
        node.left = self.null_node
        node.right = self.null_node
        node.color = 1
        y = None
        x = self.root
        while x != self.null_node:
            y = x
            if node.key < x.key:
                x = x.left
            else:
                x = x.right
        node.parent = y
        if y == None:
            self.root = node
        elif node.key < y.key:
            y.left = node
        else:
            y.right = node
        if node.parent == None:
            node.color = 0
            return
        if node.parent.parent == None:
            return
        self.fix_insert(node)

    def fix_insert(self, k):
        while k.parent.color == 1:
            if k.parent == k.parent.parent.right:
                u = k.parent.parent.left
                if u.color == 1:
                    u.color = 0
                    k.parent.color = 0
                    k.parent.parent.color = 1
                    k = k.parent.parent
                else:
                    if k == k.parent.left:
                        k = k.parent
                        self.right_rotate(k)
                    k.parent.color = 0
                    k.parent.parent.color = 1
                    self.left_rotate(k.parent.parent)
            else:
                u = k.parent.parent.right

                if u.color == 1:
                    u.color = 0
                    k.parent.color = 0
                    k.parent.parent.color = 1
                    k = k.parent.parent
                else:
                    if k == k.parent.right:
                        k = k.parent
                        self.left_rotate(k)
                    k.parent.color = 0
                    k.parent.parent.color = 1
                    self.right_rotate(k.parent.parent)
            if k == self.root:
                break
        self.root.color = 0

    def get_root(self):
        return self.root

    def search_tree(self, node, key):
        if node == self.null_node or key == node.key:
            return node
        if key < node.key:
            return self.search_tree(node.left, key)
        return self.search_tree(node.right, key)

    def find_kth_smallest(self, node, k):
        left_size = node.left.size if node.left else 0
        if k <= left_size:
            return self.find_kth_smallest(node.left, k)
        elif k == left_size + 1:
            return node
        else:
            return self.find_kth_smallest(node.right, k - left_size - 1)

    def size(self, node):
        if node == self.null_node:
            return 0
        return 1 + self.size(node.left) + self.size(node.right)

    def insert_node(self, node):
        y = None
        x = self.root
        while x != self.null_node:
            y = x
            if node.key < x.key:
                x = x.left
            else:
                x = x.right
        node.parent = y
        if y == None:
            self.root = node
        elif node.key < y.key:
            y.left = node
        else:
            y.right = node
        node.left = self.null_node
        node.right = self.null_node
        node.color = 1
        if node.parent == None:
            node.color = 0
            return
        if node.parent.parent == None:
            return
        self.fix_insert(node)

# 扩展：红黑树的删除操作
    def delete_node(self, node):
        if node == self.null_node:
            return
        y = node
        y_original_color = y.color
        if node.left == self.null_node:
            x = node.right
            self.transplant(node, node.right)
        elif node.right == self.null_node:
            x = node.left
            self.transplant(node, node.left)
        else:
            y = self.minimum(node.right)
            y_original_color = y.color
            x = y.right
            if y.parent == node:
                x.parent = y
            else:
                self.transplant(y, y.right)
                y.right = node.right
                y.right.parent = y
            self.transplant(node, y)
            y.left = node.left
            y.left.parent = y
            y.color = node.color
        if y_original_color == 0:
            self.fix_delete(x)

    def fix_delete(self, x):
        while x != self.root and x.color == 0:
            if x == x.parent.left:
                s = x.parent.right
                if s.color == 1:
                    s.color = 0
                    x.parent.color = 1
                    self.left_rotate(x.parent)
                    s = x.parent.right
                if s.left.color == 0 and s.right.color == 0:
                    s.color = 1
                    x = x.parent
                else:
                    if s.right.color == 0:
                        s.left.color = 0
                        s.color = 1
                        self.right_rotate(s)
                        s = x.parent.right
                    s.color = x.parent.color
                    x.parent.color = 0
                    s.right.color = 0
                    self.left_rotate(x.parent)
                    x = self.root
            else:
                s = x.parent.left
                if s.color == 1:
                    s.color = 0
                    x.parent.color = 1
                    self.right_rotate(x.parent)
                    s = x.parent.left
                if s.right.color == 0 and s.right.color == 0:
                    s.color = 1
                    x = x.parent
                else:
                    if s.left.color == 0:
                        s.right.color = 0
                        s.color = 1
                        self.left_rotate(s)
                        s = x.parent.left
                    s.color = x.parent.color
                    x.parent.color = 0
                    s.left.color = 0
                    self.right_rotate(x.parent)
                    x = self.root
        x.color = 0

    def transplant(self, u, v):
        if u.parent == None:
            self.root = v
        elif u == u.parent.left:
            u.parent.left = v
        else:
            u.parent.right = v
        v.parent = u.parent

    def minimum(self, node):
        while node.left != self.null_node:
            node = node.left
        return node

# 示例
tree = RedBlackTree()
tree.insert_node(Node(10))
tree.insert_node(Node(20))
tree.insert_node(Node(30))
tree.insert_node(Node(100))
tree.insert_node(Node(90))
tree.insert_node(Node(40))
tree.insert_node(Node(50))
tree.insert_node(Node(60))
tree.insert_node(Node(70))
tree.insert_node(Node(80))
print(tree.find_kth_smallest(tree.get_root(), 3).key) # 输出30

--相关问题--: