Tubetrue01
0%

从 TreeMap 源码剖析红黑树（中）插入

发表于 更新于

引言

上篇文章介绍了红黑树的性质，本篇文章将继续讲解红黑树的插入操作。

基本属性

这里我们先把红黑树的大体属性以 TreeMap 为原型列举出来。

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
public class RBTree<K extends Comparable<K>, V> {
    /*
     * 红黑树的根节点
     */
    private Entry<K, V> root;
    /*
     * 这里用 boolean 值代表红色或者黑色
     */
    private static final boolean RED   = false;
    private static final boolean BLACK = true;

    /**
     * 红黑树节点
     */
    private static class Entry<K, V> {
        private K k;
        private V v;
        /*
         * 红黑树节点是双向指针，这里指向当前节点的父节点
         */
        private Entry<K,V> parent;
        /*
         * 左节点
         */
        private Entry<K,V> left;
        /*
         * 右节点
         */
        private Entry<K,V> right;
        /*
         * 初始颜色为黑色
         */ 
        boolean color = BLACK;
        
        /**
         * 添加一个构造器
         */
        public Entry(K k, V v, Entry<K, V> parent) {
            this.k = k;
            this.v = v;
            this.parent = parent;
        }
        
        /**
         * 为当前 Entry 设置新值，并且返回老值
         */
        public V setValue(V v) {
            V old = this.v;
            this.v = v;
            return old;
        }
    }
    /**
     * 获取当前节点的父节点
     */ 
    private static <K,V> Entry<K,V> parentOf(Entry<K,V> p) {
      return (p == null ? null: p.parent);
    }
        
    /**
     * 获取当前节点的左子节点
     */ 
    private static <K,V> Entry<K,V> leftOf(Entry<K,V> p) {
      return (p == null) ? null: p.left;
    }
   
    /**
     * 获取当前节点的右子节点
     */  
    private static <K,V> Entry<K,V> rightOf(Entry<K,V> p) {
      return (p == null) ? null: p.right;
    }
    
    /**
     * 获取当前节点的颜色，这里有一点需要注意，就是当 p 为 null 时，
     * 返回黑色，换句话说，黑色作为了一种判断节点为空的方式。
     */
    private static <K,V> boolean colorOf(Entry<K,V> p) {
      return (p == null ? BLACK : p.color);
    }
}

常规操作

变色

节点的颜色由黑变红或者由红变黑。

1
2
3
4
5
6
7
/**
 * 方法比较简单，就是将节点 p 设置成红色或者黑色（true：黑色，false：红色）
 */
private static <K,V> void setColor(Entry<K,V> p, boolean c) {
    if (p != null)
        p.color = c;
}

左旋

以某个节点作为旋转点，其右子节点变为旋转节点的父节点，右子节点的左子节点变为旋转节点的右子节点，左子节点保持不变。

根据上图，我们看下左旋的代码实现：

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
/**
 * 以 p 为旋转点进行左旋
 * 
 *       5 (p)                 7
 *      / \                   / \
 * (l) 4   7 (r)   -->       5   8
 *        / \               / \
 *       6   8             4   6
 */
private void rotateLeft(Entry<K, V> p) {
    if (p != null) {
        // 获取 p 的右节点
        Entry<K, V> r = p.right;
       
        // 将 p 的右节点指向 r 的左节点
        p.right = r.left;
        // 因为是双向指针，所以要重新设置 r.left 的指向
        if (r.left != null)
            r.left.parent = p;

        // 将 r 的父节点指向 p 的父节点
        r.parent = p.parent;
        
        // 如果 p 的父节点为 null，说明 p 是根节点，那么就将根节点指向 r 
        if (p.parent == null) 
            root = r;
        // 如果 p 的父节点不为空，则需要判断 p 节点是其父节点的左节点还是右节点
        else if (p.parent.left == p)
            p.parent.left = r;
        else 
            p.parent.right = r;
        // 将 r 的左节点指向 p
        r.left = p;
        // 将 r 作为 p 的父节点
        p.parent = r;

    }
}

整体逻辑比较清晰，唯一要注意的一点是，红黑树的各个节点之间是双向指针，在更改左右节点的同时还要更新其父节点的指向。

右旋

以某个节点作为旋转点，其左子节点变为旋转节点的父节点，左子节点的右子节点变为旋转节点的左子节点，右子节点保持不变。

有了左旋的基础，我们在分析右旋就简单多了，只要把左旋中的方向左右兑换就行了。但是我们为了加深理解，还是在捋一遍吧！

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
/**
 * 以 p 为旋转点进行右旋
 *
 *        7 (p)              5
 *       / \                / \
 * (l)  5   8 (r)   -->    4   7
 *     / \                    / \
 *    4   6                  6   8
 */
private void rotateRight(Entry<K, V> p){
    if (p != null) {
        Entry<K, V> l = p.left;
        
        // 将 l 的右节点挂载到 p 下
        p.left = l.right;
        if (l.right != null)
            l.right.parent = p;
        
        // 更新 l 的父节点指向
        l.parent = p.parent;
        
        // p 为空，说明是根节点，将 l 更新为根节点
        if (p.parent == null)
            root = l;
        else if (p.parent.left == p)
            p.parent.left = l;
        else 
            p.parent.right = l;
        
        l.right = p;
        p.parent = l;

    }
}

红黑树的插入

2-3-4 树的插入

要想理解红黑树的插入，我们首先要明白 2-3-4 树的插入操作。

2-3-4 树的插入一定是从叶子开始的，也就是说它是一棵自底向上增长的树。当叶子满足合并条件时，会进行合并。如果不满足条件时，中间的元素会向上增长。我们根据下图来回顾一下 2-3-4 树的增长过程：

还记得上篇文章中介绍的红黑树与 2-3-4 树的等价关系吗？我们搞懂了 2-3-4 树的插入逻辑，那么我们把红黑树的插入逻辑也以图例的形式展现出来（有一点需要特别注意，红黑树的插入节点都是红色的）：

只要我们记住红黑树与 2-3-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
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
/**
 * 我们将 4 作为带插入的元素
 *
 * (t、p) 8               8              8            8          8
 *       / \             / \            / \          / \        / \
 *      5   12 -> (t、p) 5   12   ->   5   12  ->   5   12  -> 5   12
 *     / \              / \           / \          / \        / \
 *    3   6            3   6   (t、p) 3  6     (p) 3  6       3  6
 *                                                 \          \
 *                                             (t) null        4
 */
public V put(K key, V value) {
    // 这里直接判断 key 是否为空，如果为 null 就抛出空指针异常。
    if (key == null)
        throw new NullPointerException();

    Entry<K, V> t = root;

    // 如果根节点为 null
    if (t == null) {

        // 直接作为根节点
        root = new Entry<>(key, value, null);
        return null;
    }

    // 当根节点不为 null 时，就要添加节点了
    int cmp;
    Entry<K, V> parent;

    do {
        // 先把 t 当前的值保留下来，因为在接下来遍历的过程中，t 可能为 null，此时 parent 就是待插入节点的父节点
        parent = t;
        // 从根开始遍历，找到插入的目标位置节点
        cmp = key.compareTo(t.k);
        // 如果待插入的节点比当前节点大
        if (cmp > 0)
            t = t.right;
            // 待插入的节点比当前节点小
        else if (cmp < 0)
            t = t.left;
        else
            // key 相等，则覆盖，并且返回旧值
            return t.setValue(value);
    } while (t != null);

    // 创建新的节点，并且将将该节点的父节点指向 parent
    Entry<K, V> newEntry = new Entry<>(key, value, parent);
    // 因为指针是双向的，这里要确定是 parent 的左孩子还是右孩子
    if (cmp > 0)
        parent.right = newEntry;
    else
        parent.left = newEntry;

    // 元素插入之后，我们开始调整树的结构
    fixAfterInsertion(newEntry);
    return null;
}

红黑树的插入操作就介绍完了，从代码上来看没有什么复杂性。而复杂的核心就在 fixAfterInsertion(newEntry) 方法了。不过不要着急，我们马上开始！

红黑树的调整

在讲代码之前，我们需要先弄清楚需要调整的场景，只有把逻辑捋顺了，看代码才不那么抽象。

3 节点调整

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
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
/**
 * @param x 为刚刚插入的元素，我们以它为基点进行树的调节，谁让它破坏了树的平衡呢
 *============================3 节点==================================
 *            3 (黑)     1 (黑)                  2 (黑)
 *           /            \                     / \
 *          2 (红)         2 (红)     ->   (红) 1   3 (红)
 *         /                \
 *       *1 (红)            *3 (红)
 *--------------------------------------------------------------------
 *     4 (黑)        4 (黑)            3 (黑)
 *    /             /                /  \
 *   2 (红)   ->    3 (红)  ->  (红) 2   4 (红)
 *    \            /
 *     *3 (红)     2 (红)
 *--------------------------------------------------------------------
 *     3 (黑)           3 (黑)               4 (黑)
 *      \               \                  /  \
 *       5 (红)    ->    4 (红)   ->  (红) 3    5 (红)
 *       /               \
 *     *4 (红)            5 (红)
 *============================4 节点==================================
 *            2 (黑)                2 (红)              2 (黑)
 *           / \                   / \                 / \
 *      (红) 1  3 (红)    ->  (黑) 1  3 (黑)   ->  (黑) 1   3 (黑)
 *               \                   \                    \
 *               *4 (红)              4 (红)               4 (红)
 *--------------------------------------------------------------------
 *            3 (黑)                3 (红)              3 (黑)
 *           / \                   / \                 / \
 *      (红) 1  5 (红)    ->  (黑) 1  5 (黑)   ->  (黑) 1   5 (黑)
 *              /                    /                    /
 *            *4 (红)                4 (红)               4 (红)
 *--------------------------------------------------------------------
 *            5 (黑)              5 (红)                5 (黑)
 *           / \                 / \                  / \
 *     (红) 4   6 (红)  ->  (黑) 4   6 (黑)   ->  (黑) 4  6 (黑)
 *         /                   /                     /
 *   (红) *2              (红) 2                (红) 2
 *--------------------------------------------------------------------
 *            7 (黑)              7 (红)                7 (黑)
 *           / \                 / \                  / \
 *     (红) 4   8 (红)  ->  (黑) 4   8 (黑)   ->  (黑) 4  8 (黑)
 *          \                   /                    /
 *     (红) *5             (红) 5               (红) 5
 *====================================================================
 */
private void fixAfterInsertion(Entry<K, V> x) {
    // 新插入的元素都是红色
    x.color = RED;
    /**
     * 只有当 x 不是 null，不是根节点以及为 x 的父节点为红色时才进行遍历调整。
     * 如果 x 为根，直接将颜色改为黑色即可；如果 x 的父节点颜色为红色时，说明出现了红色相连，此时才需要调整。其实此时只涉及
     * 了 3 节点以及 4 节点的场景。因为 2 节点可以直接合并，毕竟 2 节点时，
     * 它的父颜色是黑色。
     */
    while (x != null && x != root && x.parent.color == RED) {
        // 如果是左倾树
        if (parentOf(x) == leftOf(parentOf(parentOf(x)))) {
            Entry<K, V> y = rightOf(parentOf(parentOf(x)));
            // 如果 x 的祖父节点存在右节点，说明当前是 4 节点，那么只需要变色就行了
            if (colorOf(y) == RED) {
                // 将 x 的父节点、父的兄弟节点设置为黑色，祖父节点设置为红色
                setColor(parentOf(x), BLACK);
                setColor(parentOf(parentOf(x)), RED);
                setColor(y, BLACK);
                // 由于祖父节点被设置了红色，所以还要将祖父节点作为下一次迭代的基点进行调整
                x = parentOf(parentOf(x));
            } else {
                // y 的颜色为黑色，说明 x 的祖父没有右节点，即：当前是 3 节点，那么就需要变色加旋转了
                // 在进行大的旋转之前，先判断是否需要进行小的旋转，也就是说是否 x 是作为右倾树出现的，那么先转换为左倾树
                if (x == rightOf(parentOf(x))) {
                    // 旋转之后，x 会跟 x 的父调换位置，所以需要将父赋给 x，这样旋转之后，最底层的还是 x
                    x = parentOf(x);
                    rotateLeft(x);
                }
                // 3 节点，父变黑，祖父变红，然后以祖父为中心右旋
                setColor(parentOf(x), BLACK);
                setColor(parentOf(parentOf(x)), RED);
                rotateRight(parentOf(parentOf(x)));
            }

        } else {
            // 右倾树
            // 先拿到祖父的左节点，以此来判断是 3 节点还是 4 节点
            Entry<K, V> y = leftOf(parentOf(parentOf(x)));
            if (colorOf(y) == RED) {
                // 存在，说明是 4 节点，只需要单穿的染色就行了
                setColor(parentOf(x), BLACK);
                setColor(parentOf(parentOf(x)), RED);
                setColor(y, BLACK);
                // 既然祖父已经是红色了，那么肯定需要再以祖父为中心进行调整了，万一它的父节点也是红色的呢
                x = parentOf(parentOf(x));
            } else {
                // 不用说，肯定是 3 节点了
                // 如果是左倾树，那么需要先旋转一次变成右倾树，因为整个子树是右倾的
                if (x == leftOf(parentOf(x))) {
                    x = parentOf(x);
                    rotateRight(x);
                }
                setColor(parentOf(x), BLACK);
                setColor(parentOf(parentOf(x)), RED);
                rotateLeft(parentOf(parentOf(x)));
            }

        }
    }

    // 最后将根设置为黑色
    root.color = BLACK;
}

根据注释中的各种场景，在理解代码就简单很多了。

尾声

至此，红黑树的插入操作就介绍完了。整个插入过程还是比较简单的，通过迭代确定好插入的位置，然后进行指向。难点其实在于插入之后的平衡调节，相比 AVL 树，由于红黑树属于黑色平衡的弱平衡，其实只有在 3 节点的场景中才会涉及到旋转。4 节点的分裂也只是单纯的染色而已。另外需要注意的一点是，红黑树所有的插入节点一定是红色，同时节点之间是双向指针。在旋转变换时，一定要记得修改！

下一篇文章就是红黑树系列的最后一篇【删除】了，也是整个红黑树中最难的部分，不过不要紧，我还是以最简单的方式为大家呈现出来。好，我们下篇见！