【ALGORITHM.0x04】从二叉搜索树开始手撕红黑树

本文最后更新于：2022年5月30日 下午

面试官：听说你算法还行，来给👴简单手写个红黑树就让你过了

0x00.一切开始之前

红黑树（red black tree）是自平衡二叉搜索树的一种，其特点在于能够高效地在插入与删除操作时实现自平衡，可以在 O(log N) 时间内完成查找、插入、删除，非常恐怖，相应地红黑树的实现代码也非常复杂，因此据传某大厂面试官为了刁难面试者就会天天考他们手撕红黑树

刚好笔者在阅读 Linux kernel 源码的过程中碰到不少使用到红黑树这一数据结构的地方（比如说 vma、task_struct…），此前也未深究过这一数据结构，所以今天来简单讲一讲如何从零开始手撕一棵红黑树；）

会参考《算法导论（第三版）》，但是不会照抄上面的代码，主要还是借鉴思路；）

当然，《算法导论》上面也只有伪代码

Pre.二叉树节点定义

这里还是用笔者最熟悉的 C 语言来定义一个二叉树：

1
2
3
4
5
6
struct BinaryTreeNode
{
    size_t key;
    struct BinaryTreeNode *left;
    struct BinaryTreeNode *right;
};

我们将一棵树定义如下，本质上就是根节点的 wrapper，不过这样可以方便我们更改根节点：

1
2
3
4
struct Tree
{
    struct BinaryTreeNode *root;
};

0x01. 二叉搜索树

我们先从二叉搜索树的基本概念讲起，顾名思义，一棵二叉搜索树（binary search tree）本质上是一棵二叉树，不过其多了额外的性质：

  • 设 x 为二叉搜索树中的一个节点：
    • 若 y 为 x 的左子树中的一个非空节点，则有 y.key <= x.key
    • 若 y 为 x 的右子树中的一个非空节点，则有 y.key >= x.key

即一个节点的左子树中节点的值总小于等于该节点，右子树中节点的值总大于等于该节点，这便是二叉搜索树

查找节点

查找基本上没有什么好讲的点，后面查找基本上也都是这个思路，所以这里讲完查找之后后面不会再重复叙述

为什么叫搜索树呢？因为有了这个性质我们便很容易在整棵树当中寻找特定值的节点，例如从根节点开始搜索：

递归查找

  • 若当前节点为空或当前节点的值等于被搜寻值，返回结果
  • 若当前节点的值大于被搜寻值，搜索当前节点的左子树
  • 若当前节点的值小于被搜寻值，搜索当前节点的右子树

有了以上算法我们很容易就能写出递归查找某节点的代码：

1
2
3
4
5
6
7
8
9
struct BinaryTreeNode *searchNode(struct BinaryTreeNode *root, size_t dst_val)
{
    if (!root || root->key == dst_val)
        return root;
    else if (root->key > dst_val)
        return searchNode(root->left, dst_val);
    else
        return searchNode(root->right, dst_val);
}

递推查找

当然，若是树的深度太高，递归很容易爆栈（例如内核这种栈只有一两张内存页的场景），因此接下来我们将递归变递推：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
struct BinaryTreeNode *searchNode(struct BinaryTreeNode *root, size_t dst_val)
{
    while (root)
    {
        if (root->key == dst_val)
            break;
        else if (root->key > dst_val)
            root = root->left;
        else
            root = root->right;
    }
    
    return root;
}

插入节点

我们在插入节点时需要保证二叉搜索树依然是二叉搜索树，因此我们在插入时遵循如下算法，从根节点开始遍历：

  • 若待插入节点的值小于当前节点的值
    • 若左子树为空，则待插入节点成为当前节点的左孩子
    • 若左子树非空，则继续插入到左子树中
  • 若待插入节点的值大于当前节点的值
    • 若右子树为空，则待插入节点成为当前节点的右孩子
    • 若右子树非空，则继续插入到右子树中
  • 若待插入节点的值等于当前节点的值
    • 自定义操作

若出现待插入节点的值在这棵二叉搜索树中已经存在的情况，二叉搜索树对于这种情况并没有严格的定义，你可以选择为当前节点增加一个 count 成员然后++，亦或是什么也不做

递归插入

这里我们不直接用一个节点来插入，而是通过一层 wrapper 来进行插入（参见 pre 部分）

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
#define NODE_EXISTS -1

int insertNode(struct Tree *tree, struct BinaryTreeNode *node)
{
    if (!tree->root)
    {
        tree->root = node;
        return 0;
    }
    
    return insertNodeRecursion(tree->root, node);
}

int insertNodeRecursion(struct BinaryTreeNode *root, struct BinaryTreeNode *node)
{
    if (root->key < node->key)
    {
        if (!root->right)
        {
            root->right = node;
            return 0;
        }
        else 
        {
            return insertNode(root->right, node);
        }
    }
    else if (root->key > node->key)
    {
        if (!root->left)
        {
            root->left = node;
            return 0;
        }
        else 
        {
            return insertNode(root->left, node);
        }
    }
    else
    {
        // do nothing
        return NODE_EXISTS;
    }
}

递推插入

魔改的《算法导论》上的算法，递推遍历找到待插入位置，在遍历的过程中 parent 节点其实就是 child 的父节点，当 root 为 NULL 时说明找到了可插入位置，于是把待插入节点插到 parent 上

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
#define NODE_EXISTS -1

int insertNode(struct Tree *tree, struct BinaryTreeNode *node)
{
    struct BinaryTreeNode *parent;
    struct BinaryTreeNode *child;
    
    child = tree->root;
    parent = NULL;
    while (child)
    {
        parent = child;
        if (node->key < child->key)
            child = child->left;
        else if (node->key > child->key)
            child = child->right;
        else
            return NODE_EXISTS;
    }
    
    // tree is mpty
    if (!parent)
        tree->root = node;
    else if (node->key < parent->key)
        parent->left = node;
    else if (node->key > parent->key)
        parent->right = node;
    else
        return NODE_EXISTS;
    
    return 0;
}

删除节点

我们需要在删除掉指定节点后保持这棵树仍是一棵二叉搜索树，因此删除节点比插入节点要稍微麻烦一点点

删除节点主要存在以下三种情况：

  • 待删除节点 x 没有子节点
    • 直接删除，父节点对应子树置NULL
  • 待删除节点 x 有一个子节点 y
    • 让 y 替代 x 原来的位置
  • 待删除节点 x 有两个子节点
    • 寻找 x 的后继节点 y，让 y 替代 x 原来的位置，x 原来的左右子树成为 y 的新的左右子树

后继节点：二叉树中序遍历中当前节点的后一个节点

第三种情况是比较复杂的，因为我们同时还需要考虑到 y 的位置，若 y 是 x 的右孩子，则直接让 y 上位替代 x 即可；若 y 不是 x 的右孩子，则 y 上位替代 x 之后由 y 的右孩子替代 y 的原位置（作为 z 的后继节点，y 必定没有左孩子）

在《算法导论》中最终将这三种情况按如下算法路径进行处理，对于待删除结点 x 而言：

  • 若 x 没有左孩子，则用其右孩子来替换 x（这个右孩子可以是 NULL，所以实际包含两种情况：仅有一右孩|无子结点）
  • 若 x 仅有一个孩子且为左孩子，则用其左孩子来替换 x
  • 若 x 同时有左右孩子，查找 x 的后继结点 y，将 y 移出原来的位置进行拼接，并替换树中的 x
    • 若 y 为 x 的右孩子，则直接用 y 替换 x
    • 若 y 为 x 的右子树中结点，先用 y 的右孩子替换 y，之后用 y 替换 x

算法导论上的图

最后就是删除操作的代码实现了，这里注意后继节点的一些性质：

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
#define EMPTY_TREE -1

struct BinaryTreeNode *findParent(struct Tree *tree, struct BinaryTreeNode *node)
{
    struct BinaryTreeNode *root;
    
    if (!tree || tree->root == node)
        return NULL;
    
    root = tree->root;
    while (root)
    {
        if (root->left == node || root->right == node)
            return root;
        if (root->key < node->key)
            root = root->right;
        else
            root = root->left;
    }
    
    return NULL;
}

struct BinaryTreeNode *findSuccessor(struct BinaryTreeNode *node)
{
    struct BinaryTreeNode *successor;
    
    successor = node->right;
    while (successor && successor->left)
        successor = successor->left;
    
    return successor;
}

int deleteNode(struct Tree *tree, struct BinaryTreeNode *dst)
{
    struct BinaryTreeNode *successor;
    struct BinaryTreeNode *successor_parent;
    struct BinaryTreeNode *dst_parent;
    struct BinaryTreeNode *replayer;
    
    if (!tree || !tree->root)
        return EMPTY_TREE;
    
    dst_parent = findParent(tree, dst);
    
    if (!dst->left) // no left child, right to replace x
    {
        replayer = dst->right;
    }
    else if (!dst->right) // only left child, left to replace x
    {
        replayer = dst->left;
    }
    else // both children exist
    {
        successor = findSuccessor(dst);
        successor_parent = findParent(tree, successor);

        if (dst != successor_parent) // successor in dst's right tree
        {
            successor_parent->left = successor->right;
            successor->right = dst->right;
        }
        
        successor->left = dst->left;
        replayer = successor;
    }
    
    // replace x
    if (!dst_parent)
        tree->root = replayer;
    else if (dst == dst_parent->left)
        dst_parent->left = replayer;
    else
        dst_parent->right = replayer;
    
     return 0;
}

0x02. AVL树

AVL树（Adelson-Velsky and Landis Tree）是一种自平衡二叉查找树，其会在插入与删除节点的过程中保持自身的一种平衡，由于这个性质，其也被称为 高度平衡树 ，其性质为：

  • AVL 树任一节点对应的两棵子树的最大高度差为 1

AVL 树最初是被发明来解决二叉查找树在经过多次操作后出现的不平衡的情况，例如下图这种情况，虽然其仍是一棵二叉查找树，但是由于其结构的不平衡，导致查找的效率大幅降低

image.png

比如说你要查找最左边的那一部分节点，这个时候二叉搜索树就退化成一个数组了，时间复杂度直接增长到 O(N)，这显然不是我们想要看到的

AVL 树 通过在增加与删除元素时进行一次或多次旋转操作来实现树的重新平衡，其查找、插入、删除操作在平均与最坏情况下的时间复杂度都是 O(logN)

Pre. AVL 树节点定义

这里我们为节点引入一个新的成员——高度：由根节点到其左右子树的【最大】深度

高度是我们对树进行重构使其重新成为 AVL 树的依据，我们在每个节点中记录下该节点的高度，并在每次操作时进行动态更新

1
2
3
4
5
6
7
struct ALVTreeNode
{
    size_t	key;
    size_t	height;
    struct ALVTreeNode 	*left;
    struct ALVTreeNode	*right;
}

相应地，我们应当有一个用以计算一棵树的高度的方法：

1

旋转操作

这里我们引入一个新的操作——旋转（rotation），这是组成我们接下来的插入/删除操作中的一个子操作，我们通过单次/多次旋转操作来保持 AVL 树的高度平衡的特性

下图是一个最简单的旋转的例子，通过旋转使得一棵树重新成为 AVL 树

260px-AVL_Tree_Example.gif

左旋

右旋

插入节点

删除节点

0x03. 2-3-4树

0x04. 红黑树