前言

红黑树真是又臭又长，从 STL 看到红黑树到今天基本理解与实现花了快三天时间，相当于写 vector 与 list 的时间总和，令人感叹。

红黑树性质

  1. 节点是红色或黑色。
  2. 根是黑色。（根据OI_WIKI，这一条不一定需要满足）
  3. 所有叶子都是黑色（叶子是 NIL 节点）。
  4. 每个红色节点必须有两个黑色的子节点。（从每个叶子到根的所有路径上不能有两个连续的红色节点。）
  5. 从任一节点到其每个叶子的所有简单路径都包含相同数目的黑色节点。

红黑树结构

节点结构

1
2
3
4
5
6
7
8
class Node():
    """0 is black, 1 is red"""
    def __init__(self, val):
        self.val = val
        self.parent = None
        self.left = None
        self.right = None
        self.color = 1

红黑树结构

需要注意的是这里的 TNULL 设计，代表了红黑树中的叶子节点（空节点），但是与 None 不同的是，其具有 color 与 parent 等属性，这样设计可以使后续的许多判断不用单独考虑 None 的情况。

1
2
3
4
5
6
7
class RedBlackTree():
    def __init__(self):
        self.TNULL = Node(0)
        self.TNULL.color = 0
        self.TNULL.left = None
        self.TNULL.right = None
        self.root = self.TNULL

左旋与右旋

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
def left_rotate(self, node):
    '''
    左旋操作，分为三步：
    1. 将 r.parent 置为 p（可能会更新 root），根据 node 的位置把 r 连接至 p.left 或 p.right
    2. 将 node.parent 置为 r，r.left 置为 node
    3. 将 node.right 置为 rl，若 rl 存在，将 rl.parent 置为 node，
    * 左旋示意图：对节点x进行左旋
    *     p                       p
    *    /                       /
    *   node                    r
    *  / \                     / \
    * x   r     ----->      node  rr
    *    / \                /  \
    *   rl  rr             x   rl
    '''

    p = node.parent 
    r = node.right
    rl = r.left

    # 1 r 连到父节点，若父节点为空（node 为 root），更新 root
    r.parent = p
    if not p: self.root = r 
    else:
        # 1.1 根据 node 相对于 p 的位置来连接 p 与 r
        if (p.left == node): p.left = r 
        else: p.right = r 

    # 2. 连接 node 与 r
    node.parent = r
    r.left = node 
    
    # 3. 连接 rl 与 node
    node.right = rl 
    if (rl): rl.parent = node 

def right_rotate(self, node):
    ''' 
    右旋操作，分为三步：
    1. 将 l.parent 置为 p（可能会更新 root），根据 node 的位置把 l 连接至 p.left 或 p.right
    2. 将 node.parent 置为 l，l.right 置为 node
    3. 将 node.left 置为 lr，若 lr 存在，将 lr.parent 置为 node，
    * 右旋示意图：对节点y进行右旋
    *         p                   p
    *        /                   /
    *      node                 l
    *     /    \               / \
    *    l      r   ----->   ll  node
    *   / \                      /  \
    * ll   lr                   lr   r
    '''
    p = node.parent
    l = node.left
    lr = l.right

    l.parent = p
    if not p: self.root = l
    else:
        if (p.left == node): p.left = l 
        else: p.right = l 

    node.parent = l 
    l.right = node 

    node.left = lr
    if (lr): lr.parent = node

节点插入及重新平衡

红黑树插入节点有两个步骤，节点插入以及插入新节点后的红黑树重新平衡。

节点插入

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
def insert_node(self, key):
    node = Node(key)
    node.parent = None
    node.left = self.TNULL  # 将左右节点都设置为 TNULL，即叶子节点
    node.right = self.TNULL
    node.color = 1          # 新插入的节点为红色

    y = None
    x = self.root

    # 从根节点开始，找到新插入节点的位置
    while x != self.TNULL:
        y = x
        if node.val < x.val: x = x.left
        else: x = x.right

    # 插入新节点
    node.parent = y
    if y == None: self.root = node
    elif node.val < y.val: y.left = node
    else: y.right = node

    self.insert_fix(node)  # 修正红黑树

插入后修正

在讨论插入情况之前，先做如下表述上的约定：新插入节点为 N (Node)，其父节点为 P (parent)， 祖父节点为 G (Grandparent)，叔叔节点为 U (Uncle)，插入节点与父节点的相对位置为 L / R。

新插入节点初始化为红色节点，插入节点时有以下几种情况：

  1. 插入节点为根节点，直接将节点颜色置为黑色
  2. 插入节点的父节点为黑色，不需要做任何操作
  3. 插入节点的父节点为红色，需要进行调整

    1. 插入节点的叔叔节点为红色，将父节点和叔叔节点置为黑色，祖父节点置为红色，将祖父节点作为新的插入节点，继续调整（因为此时祖父节点可能也不满足性质 4）。

    2. 插入节点的叔叔节点为黑色，且插入节点与父节点同侧（LL/RR），以祖父节点为支点进行反向旋转（R/L），将父节点置为黑色，祖父节点置为红色。

    3. 插入节点的叔叔节点为黑色，且插入节点与父节点异侧（LR/RL），以父节点为支点进行反向旋转（L/R），转化为 3.2。

可以看出，由于将新节点都设置为了红色节点，因此插入操作有可能违背的就只有性质4，即红色节点的父节点不能为红色。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
def insert_fix(self, node):
    # 父节点为黑色不用处理，case2
    while (node.parent and node.parent.color == 1):  
        # L
        if (node.parent == node.parent.parent.left): 
            u = node.parent.parent.right
            # uncle 红色，case3.1
            # 由于设置了 TNULL，因此不用判断 u 是否为 None
            if (u.color == 1):    
                node.parent.color = u.color = 0
                node.parent.parent.color = 1
                # 更新当前处理节点切换为祖父节点
                node = node.parent.parent
            else:
                # LR, case3.3 转化为 case3.2
                if (node == node.parent.right):
                    node = node.parent
                    self.left_rotate(node)
                node.parent.color = 0
                node.parent.parent.color = 1
                self.right_rotate(node.parent.parent)
        # R
        else:
            u = node.parent.parent.left
            if (u.color == 1):
                node.parent.color = u.color = 0
                node.parent.parent.color = 1
                node = node.parent.parent
            else:
                # RL, 转为 case3.2
                if (node == node.parent.left): 
                    node = node.parent
                    self.right_rotate(node)
                node.parent.color = 0
                node.parent.parent.color = 1
                self.left_rotate(node.parent.parent)
        
        if (node == self.root): break
    # 根节点始终为黑色
    self.root.color = 0

在 insert_fix 的实现过程中需要注意的是，旋转会改变树的结构，因此会改变原始的父子关系，一定要利用 node.parent.parent 这样的方式来访问父亲及爷爷节点而不能提前储存后利用变量访问，这样会导致很多问题（血泪教训）

节点删除及树修复

红黑树的节点删除操作逻辑上分为两部分，节点的删除以及树的重新平衡，实际上在许多代码中也将删除操作分为两个函数来完成，这里的实现也按照此逻辑。

博主认为，网络上包括书籍中的很多关于红黑树删除的教程，之所以让人看了云里雾里，是因为它们在演示中并没有将删除以及重新平衡这两个步骤分开来展示，在列举各种情况时总是把删除以及重新平衡的情况放一起说完，这样就会导致逻辑的混乱，且不利于记忆写代码的步骤。

节点删除

就节点删除这一个步骤来看，红黑树的节点删除操作与BST（二叉搜索树）基本相同，可分为三种情况：

  1. 删除只有一个子树的节点：删除原节点，用单独的子节点代替被删除节点。
  2. 删除没有子树的节点：直接删除原节点，相当于用一个空节点 (NIL) 代替了被删除节点。
  3. 删除有两个子树的节点：寻找后继或前继节点（这里选择后继），将待删除节点的值变为后继节点的值，再删除后继节点，删除后继节点一定会变成 1 或 2 的情况（可以思考一下为什么）。

但由于红黑树的特殊性（多了一个颜色属性），因此要根据被删除节点的颜色，采取不同的后续平衡步骤中。

辅助函数

以下为删除节点所需要的两个辅助函数。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
def findMin(self, node):
    """
    找到以 node 节点为根节点的树的最小值节点 并返回
    """
    temp_node = node
    # 判断条件也要为 TNULL
    while temp_node.left != self.TNULL: temp_node = temp_node.left
    return temp_node

def transplant(self, node1, node2):
    """
    辅助函数，用 node2 替换 node1
    """
    p1 = node1.parent
    if not p1: self.root = node2
    elif (node1 == p1.left): p1.left = node2
    elif (node1 == p1.right): p1.right = node2
    # 由于 node2 必定为 TNULL，不会为 None，因此一定有 parent
    node2.parent = p1

transplant 这个函数并非简单地用 node2 替换 node1。仔细观察就会发现，这里只考虑了 node1 的 parent 指针与 node2 的交接而没有考虑两个节点的 left 及 right 指针的问题。因为在这个函数的大部分使用场景中，node2 为 node1 唯一的子节点。因此该函数的作用应该解释为 “用 node1 的唯一子节点 node2 替换自己的位置”。

节点删除

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
def delete_node(self, node, key):
    
    z = self.TNULL
    # 找到 key 对应的节点
    while node != self.TNULL:
        if node.val == key: z = node
        if node.val <= key: node = node.right
        else: node = node.left
    # 如果没有找到 key 对应的节点，返回
    if z == self.TNULL:
        print("Cannot find key in the tree")
        return
    
    node = z
    node_color = node.color
    temp_node = None  # 替换到被删除节点（node）处的节点

    # ===================== case1 与 case2 ================== //
    # node 本身被删除
    if node.left == self.TNULL:
        temp_node = node.right
        self.transplant(node, node.right)
    elif node.right == self.TNULL:
        temp_node = node.left
        self.transplant(node, node.left)
    
    # ========================= case 3 ====================== //
    # node 的后继被删除
    else:
        # 找到后继节点，及右子树中的最小节点，该节点将代替 node 被删除
        successor = self.findMin(node.right)
        node_color = successor.color
        # temp_node 将在 successor 被删除后代替它
        temp_node = successor.right

        # 另一种实现方式，但感觉较为复杂，不如直接交换两个值
        # if successor.parent == node:
        #     temp_node.parent = successor
        # else:
        #     self.transplant(successor, successor.right)
        #     successor.right = node.right
        #     node.right.parent = successor

        # self.transplant(node, successor)
        # successor.left = node.left
        # successor.left.parent = successor
        # successor.color = node.color
        
        # 将 node 替换为 successor，相当于仅替换值
        node.val = successor.val
        # 删除 successor，用 temp_node 替代
        self.transplant(successor, temp_node)
        
    # 这里的 delete_rebalance 为红黑树的重新平衡操作，需要注意的是，
    # 重新平衡的操作总是在被删除节点的位置进行的，即上面维护的 temp_node 最终的位置
    # 仅在被删除节点的颜色（node_color 维护）为黑色时才需重新平衡树
    if (node_color == 0): self.delete_fix(temp_node)

关于替换节点的两种实现方式，在 python 中直接交换两个值是没什么问题的，因为 python 有 GC 机制，断开连接的节点过一段时间就会被清除。但是 C++ 中还是需要保证与树断开连接的节点确实是我们希望删除的那个，因为还需要通过对应的指针对该节点进行析构和内存回收等操作，如果使用了值替换则会破坏树的结构；
因此，STL 中的做法也与上述注释掉的做法相同，将 successor 换到被删除的节点的位置，颜色也一同变化，再手动地析构 node。

红黑树修复

在讨论红黑树的重新平衡之前，回顾一下在红黑树的节点删除部分我们做了什么：首先，根据将要被删除的节点的孩子数量，来决定删除策略，其中零或一个孩子的情况较好处理，而两个孩子的情况则可以通过寻找后继节点的方式转换为前两种情况，需要注意的是，前两种情况删除的 node 位置即为最初的 node 位置，而后一种情况删除的节点位置已经不在最初的 node 的位置了，而是其后继所在的位置。

在删除完节点后，temp_node 所指向的节点为 替代被删除节点的节点，它的位置在被删除节点的位置，后续的所有的修复情况均以 temp_node 的位置及其相关节点的颜色作为判断标准。

其次，不难注意到，只有当被删除节点颜色为黑色时，才会增加重新平衡红黑树的操作，否则红黑树不需要调整，因此。进入 delete_fix 函数时，红黑树的情况为：所有通过 temp_node 的路径均少了一个黑色节点（因为原本 temp_node 所在位置为一个黑色节点，现在被删了）。因此，delete_fix 的目标（作用）即为：为通过 temp_node 节点的路径增加一个黑色节点，同时保证其他路径的黑色节点数目不变，才能重新平衡红黑树。

注意，这里的新增指的是通过将红色节点变为黑色节点的方式来增加黑色节点数目，而不是插入新的节点。

弄清楚了这些，才能使我们更加容易理解红黑树修复的各种情况。在下面各种情况的演示图中也可以留意每条路径上的黑色节点数量在红黑树修复前后的变化，可以更好地理解 delete_fix 的作用。

与插入操作中的定义类似，这里以 N 代表 temp_node，S 代表 N 的兄弟节点，P 代表 N 与 S 的父节点，SL 与 SR 分别代表 S 的左右子节点，节点 S 的同侧儿子表示的情况为：若 S 为 P 的左儿子，则 S 的左儿子为 S 的同侧儿子 ( SL )，反之 SR 为同侧儿子，异侧儿子定义类似。

红黑树修复有以下六种情况：

  1. N 为新的根节点

    这说明原始被删除节点即为根节点，这种情况相当于所有路径均少了一个黑色的节点，我们将 N 变为黑色，即完成操作。

  2. S 为红色

    3.png

    这种情况下无法直接处理，需要在 P 上做旋转，转换为其他情况。

  3. P、S、SL、SR 均为黑色

    这种情况下无法在这个局部结构中新增一个黑色节点（牢记我们的目标，在 N 的路径上增加一个黑色节点），此时只能将 S 变为红色，这样会使得局部结构达到平衡（左右都少了一个黑色节点），但红黑树整体依然不平衡。需要继续向上调整，试图在更上层的结构中为这条路径增加一个黑色节点。

    因此，以 P 作为 delete_fix 的参数，改为处理 P 节点。至于为什么要处理该节点，牢记：delete_fix 的作用为使过 temp_node 节点的路径增加一个黑色节点，同时保证其他路径的黑色节点数目不变，在上述的调整过程中，我们通过改变 S 的颜色的方式使这个局部结构达到了平衡，但相较于其他路径都少了一个黑色节点，因此需要处理 P 节点，使得通过 P 节点的路径增加一个黑色节点，以平衡红黑树。

  4. S、SL、SR 均为黑色，P 为红色

    这种情况比较好处理，简单地将 S 变为红色再将 P 变为黑色，就能使得 N 的路径上多出一个黑色节点，而其他路径黑色节点数量不变。

实际上 3 与 4 是同一类情况，并且这两种情况的处理可以简单地隐藏到判断条件中，现在觉得这里的划分有些不合理了，具体会在下面的代码中说明。

  1. S 为黑色，S 的同侧儿子为红色

    这种情况下我们旋转 P，交换 S 与 P 的颜色，并且将 S 的同侧儿子颜色变为黑色，就能使得所有通过 N 的路径均增加一个黑色节点而其他路径黑色节点数量不变。

  2. S 为黑色，S 的同侧儿子为黑色，异侧儿子为红色

    这种情况也不能直接处理，需要通过旋转 S 再交换 S 与它的新父亲颜色的方式转为 case5 的情况再进一步处理。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
def delete_fix(self, node):
    # 注意这里的循环条件: node.color == 0，当遇到 case3 的情况时会转而处理 P 节点
    # 此时 p == node，会再进行一次条件判断，当 p 为红色（1）时，便为 case4 的情况
    # 直接跳出循环，最后将 node.color 设为 0，结束调整。
    while (node != self.root and node.color == 0):
        if (node == node.parent.left):
            s = node.parent.right
            # case2：兄弟节点为红色节点，旋转父节点
            if (s.color == 1):
                node.parent.color = 1
                s.color = 0
                self.left_rotate(node.parent)
                s = node.parent.right  # 更新兄弟节点
            # case3: 兄弟节点及两个子节点均为黑色
            if (s.left.color == 0) and (s.right.color == 0):
                s.color = 1
                node = node.parent  # 向上继续处理
                # 此时如果 p 为红色会直接跳出循环
            else:
                # case5: 兄弟节点同侧的节点为黑色
                if (s.right.color == 0):
                    s.color = 1
                    s.left.color = 0
                    self.right_rotate(s)
                    s = node.parent.right  # 更新兄弟节点
                # 处理之后变为 case6
                # case6: 兄弟节点的同侧节点为红色
                s.color = node.parent.color
                node.parent.color = 0
                s.right.color = 0
                self.left_rotate(node.parent)
                node = self.root  # 相当于 break
        # 对称操作
        else:
            s = node.parent.left
            if (s.color == 1):
                node.parent.color = 1
                s.color = 0
                self.right_rotate(node.parent)
                s = node.parent.left
            if (s.left.color == 0) and (s.right.color == 0):
                s.color = 1
                node = node.parent
            else:
                if (s.left.color == 0):
                    s.color = 1
                    s.right.color = 0
                    self.left_rotate(s)
                    s = node.parent.left
                s.color = node.parent.color
                node.parent.color = 0
                s.left.color = 0
                self.right_rotate(node.parent)
                node = self.root  # 相当于 break
    node.color = 0  # 将 node 颜色改为黑色，可能是 root 也有可能是 case4 跳转过来的情况

辅助函数

打印出二叉树，灵感来源于力扣655题，打印二叉树

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
def printTree(self, color = True):
        
    def maxDepth(root):
        if root == self.TNULL: return -1
        leftDepth = maxDepth(root.left)
        rightDepth = maxDepth(root.right)
        return max(leftDepth, rightDepth) + 1
    
    h = maxDepth(self.root)
    m = h + 1
    n = pow(2, m) - 1
    res = [["" for i in range(n)] for i in range(m)]
    nodes = [(self.root, 0, (n-1)//2)]
    
    while nodes:
        cur, r, c = nodes.pop()
        res[r][c] = str(cur.val)
        if color: res[r][c] += "r" if cur.color == 1 else "b"
        if cur.left != self.TNULL: nodes.append((cur.left, r+1, c-pow(2,h-r-1)))
        if cur.right != self.TNULL: nodes.append((cur.right, r+1, c+pow(2,h-r-1)))
    
    return res

将一个列表插入红黑树

1
2
3
def insert_list(self, l):
    for i in l:
        self.insert_node(i)

中序遍历树

1
2
3
4
5
6
7
8
9
def tree2list(self):
    res = []
    def helper(node):
        if node is None: return
        helper(node.left)
        res.append(node.val)
        helper(node.right)
    helper(self.root)
    return res

整体测试

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
class Node():
    """0 is black, 1 is red"""
    def __init__(self, val):
        self.val = val
        self.parent = None
        self.left = None
        self.right = None
        self.color = 1


class RB_Tree():
    def __init__(self):
        self.TNULL = Node(0)
        self.TNULL.color = 0
        self.TNULL.left = None
        self.TNULL.right = None
        self.root = self.TNULL
    
    def findMin(self, node):
        """
        找到以 node 节点为根节点的树的最小值节点 并返回
        """
        while node.left != self.TNULL:
            node = node.left
        return node
        
    def transplant(self, node1, node2):
        """
        辅助函数，用 node2 替换 node1
        """
        p1 = node1.parent
        if not p1: self.root = node2
        elif (node1 == p1.left): p1.left = node2
        elif (node1 == p1.right): p1.right = node2
        # 由于 node2 必定为 TNULL，不会为 None，因此一定有 parent
        node2.parent = p1  

    def printTree(self, color = True):
        def maxDepth(root):
            if root == self.TNULL: return -1
            leftDepth = maxDepth(root.left)
            rightDepth = maxDepth(root.right)
            return max(leftDepth, rightDepth) + 1
        
        h = maxDepth(self.root)
        m = h + 1
        n = pow(2, m) - 1
        res = [["" for i in range(n)] for i in range(m)]
        nodes = [(self.root, 0, (n-1)//2)]
        
        while nodes:
            cur, r, c = nodes.pop()
            res[r][c] = str(cur.val)
            if color: res[r][c] += "r" if cur.color == 1 else "b"
            if cur.left != self.TNULL: nodes.append((cur.left, r+1, c-pow(2,h-r-1)))
            if cur.right != self.TNULL: nodes.append((cur.right, r+1, c+pow(2,h-r-1)))
        
        return res

    def left_rotate(self, node):
        '''
        左旋操作，分为三步：
        1. 将 r.parent 置为 p（可能会更新 root），根据 node 的位置把 r 连接至 p.left 或 p.right
        2. 将 node.parent 置为 r，r.left 置为 node
        3. 将 node.right 置为 rl，若 rl 存在，将 rl.parent 置为 node，
        * 左旋示意图：对节点x进行左旋
        *     p                       p
        *    /                       /
        *   node                    r
        *  / \                     / \
        * x   r     ----->      node  rr
        *    / \                /  \
        *   rl  rr             x   rl
        '''

        p = node.parent 
        r = node.right
        rl = r.left

        # 1 r 连到父节点，若父节点为空（node 为 root），更新 root
        r.parent = p
        if not p: self.root = r 
        else:
            # 1.1 根据 node 相对于 p 的位置来连接 p 与 r
            if (p.left == node): p.left = r 
            else: p.right = r 

        # 2. 连接 node 与 r
        node.parent = r
        r.left = node 
        
        # 3. 连接 rl 与 node
        node.right = rl 
        if (rl != self.TNULL): rl.parent = node 

    def right_rotate(self, node):
        ''' 
        右旋操作，分为三步：
        1. 将 l.parent 置为 p（可能会更新 root），根据 node 的位置把 l 连接至 p.left 或 p.right
        2. 将 node.parent 置为 l，l.right 置为 node
        3. 将 node.left 置为 lr，若 lr 存在，将 lr.parent 置为 node，
        * 右旋示意图：对节点y进行右旋
        *         p                   p
        *        /                   /
        *      node                 l
        *     /    \               / \
        *    l      r   ----->   ll  node
        *   / \                      /  \
        * ll   lr                   lr   r
        '''
        p = node.parent
        l = node.left
        lr = l.right

        l.parent = p
        if not p: self.root = l
        else:
            if (p.left == node): p.left = l 
            else: p.right = l 

        node.parent = l 
        l.right = node 

        node.left = lr
        if (lr != self.TNULL): lr.parent = node 

    def insert_node(self, key):
        node = Node(key)
        node.parent = None
        node.left = self.TNULL  # 将左右节点都设置为 TNULL，即叶子节点
        node.right = self.TNULL
        node.color = 1          # 新插入的节点为红色

        y = None
        x = self.root

        # 从根节点开始，找到新插入节点的位置
        while x != self.TNULL:
            y = x
            if node.val < x.val: x = x.left
            else: x = x.right

        # 插入新节点
        node.parent = y
        if y == None: self.root = node
        elif node.val < y.val: y.left = node
        else: y.right = node

        self.insert_fix(node)  # 修正红黑树
    
    def insert_fix(self, node):
        # 父节点为黑色不用处理，case2
        while (node.parent and node.parent.color == 1):  
            # L
            if (node.parent == node.parent.parent.left): 
                u = node.parent.parent.right
                # uncle 红色，case3.1
                # 由于设置了 TNULL，因此不用判断 u 是否为 None
                if (u.color == 1):    
                    node.parent.color = u.color = 0
                    node.parent.parent.color = 1
                    # 更新当前处理节点切换为祖父节点
                    node = node.parent.parent
                else:
                    # LR, case3.3 转化为 case3.2
                    if (node == node.parent.right):
                        node = node.parent
                        self.left_rotate(node)
                    node.parent.color = 0
                    node.parent.parent.color = 1
                    self.right_rotate(node.parent.parent)
            # R
            else:
                u = node.parent.parent.left
                if (u.color == 1):
                    node.parent.color = u.color = 0
                    node.parent.parent.color = 1
                    node = node.parent.parent
                else:
                    # RL, 转为 case3.2
                    if (node == node.parent.left): 
                        node = node.parent
                        self.right_rotate(node)
                    node.parent.color = 0
                    node.parent.parent.color = 1
                    self.left_rotate(node.parent.parent)
            
            if (node == self.root): break
        # 根节点始终为黑色
        self.root.color = 0
    
    def delete_node(self, key):
        z = self.TNULL
        node = self.root
        # 找到 key 对应的节点
        while node != self.TNULL:
            if node.val == key: z = node
            if node.val <= key: node = node.right
            else: node = node.left
        # 如果没有找到 key 对应的节点，返回
        if z == self.TNULL:
            print("Cannot find key in the tree")
            return
        
        node = z
        node_color = node.color
        temp_node = None  # 替换到被删除节点（node）处的节点

        # ===================== case1 与 case2 ================== //
        # node 本身被删除
        if node.left == self.TNULL:
            temp_node = node.right
            self.transplant(node, node.right)
        elif node.right == self.TNULL:
            temp_node = node.left
            self.transplant(node, node.left)
        
        # ========================= case 3 ====================== //
        # node 的后继被删除
        else:
            # 找到后继节点，及右子树中的最小节点，该节点将代替 node 被删除
            successor = self.findMin(node.right)
            node_color = successor.color
            # temp_node 将在 successor 被删除后代替它
            temp_node = successor.right

            # 另一种实现方式，但感觉较为复杂，不如直接交换两个值
            # if successor.parent == node:
            #     temp_node.parent = successor
            # else:
            #     self.transplant(successor, successor.right)
            #     successor.right = node.right
            #     node.right.parent = successor

            # self.transplant(node, successor)
            # successor.left = node.left
            # successor.left.parent = successor
            # successor.color = node.color
            
            # 将 node 替换为 successor，相当于仅替换值
            node.val = successor.val
            # 删除 successor，用 temp_node 替代
            self.transplant(successor, temp_node)

        # 这里的 delete_rebalance 为红黑树的重新平衡操作，需要注意的是，
        # 重新平衡的操作总是在被删除节点的位置进行的，即上面维护的 temp_node 最终的位置
        # 仅在被删除节点的颜色（node_color 维护）为黑色时才需重新平衡树
        if (node_color == 0): self.delete_fix(temp_node)
    
    def delete_fix(self, node):
        while (node != self.root and node.color == 0):
            if (node == node.parent.left):
                s = node.parent.right
                # case2：兄弟节点为红色节点，旋转父节点
                if (s.color == 1):
                    node.parent.color = 1
                    s.color = 0
                    self.left_rotate(node.parent)
                    s = node.parent.right  # 更新兄弟节点
                # case3: 兄弟节点的两个子节点均为黑色
                if (s.left.color == 0) and (s.right.color == 0):
                    s.color = 1
                    node = node.parent  # 向上继续处理
                else:
                    # case5: 兄弟节点同侧的节点为黑色
                    if (s.right.color == 0):
                        s.color = 1
                        s.left.color = 0
                        self.right_rotate(s)
                        s = node.parent.right  # 更新兄弟节点
                    # 处理之后变为 case6
                    # case6: 兄弟节点的同侧节点为红色
                    s.color = node.parent.color
                    node.parent.color = 0
                    s.right.color = 0
                    self.left_rotate(node.parent)
                    node = self.root  # 相当于 break
            # 对称操作
            else:
                s = node.parent.left
                if (s.color == 1):
                    node.parent.color = 1
                    s.color = 0
                    self.right_rotate(node.parent)
                    s = node.parent.left
                if (s.left.color == 0) and (s.right.color == 0):
                    s.color = 1
                    node = node.parent
                else:
                    if (s.left.color == 0):
                        s.color = 1
                        s.right.color = 0
                        self.left_rotate(s)
                        s = node.parent.left
                    s.color = node.parent.color
                    node.parent.color = 0
                    s.left.color = 0
                    self.right_rotate(node.parent)
                    node = self.root  # 相当于 break
        node.color = 0  # 将 root 设置为黑色

    def insert_list(self, l):
        for i in l:
            self.insert_node(i)
    
    def tree2list(self):
        res = []
        def helper(node):
            if node == self.TNULL: return
            helper(node.left)
            res.append(node.val)
            helper(node.right)
        helper(self.root)
        return res
1
2
3
4
import random
t = RB_Tree()
t.insert_list([ random.randint(0, 100) for i in range(10) ])
t.printTree()
[['', '', '', '', '', '', '', '94b', '', '', '', '', '', '', ''],
 ['', '', '', '59r', '', '', '', '', '', '', '', '96b', '', '', ''],
 ['', '18b', '', '', '', '81b', '', '', '', '95r', '', '', '', '100r', ''],
 ['7r', '', '26r', '', '', '', '92r', '', '', '', '', '', '', '', '']]

定义一些函数用于测试插入与删除功能

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
def insert_test(epoch = 10, num = 100):
    print("=== insert test ==")
    for _ in range(epoch):
        t = RB_Tree()
        curList = [random.randint(0, 1000) for _ in range(num)]
        t.insert_list(curList)
        curList.sort()
        ans = "pass" if curList == t.tree2list() else "fail"
        print(f"======[{ans}]======")

def remove_test(epoch = 10, num = 100):
    print("=== remove test ==")
    for _ in range(epoch):
        t = RB_Tree()
        curList = [random.randint(0, 1000) for _ in range(num)]
        t.insert_list(curList)
        ans = True
        for _ in range(random.randint(1, 10)):
            key = random.choice(curList)
            curList.remove(key)
            t.delete_node(key)
            curList.sort()
            ans = ans and (curList == t.tree2list())
        ans = "pass" if ans else "fail"
        print(f"======[{ans}]======")

def test():
    insert_test()
    remove_test()
1
test()
=== insert test ==
======[pass]======
======[pass]======
======[pass]======
======[pass]======
======[pass]======
======[pass]======
======[pass]======
======[pass]======
======[pass]======
======[pass]======
=== remove test ==
======[pass]======
======[pass]======
======[pass]======
======[pass]======
======[pass]======
======[pass]======
======[pass]======
======[pass]======
======[pass]======
======[pass]======

总结

红黑树真恶心。