Java Collections Framework 源码分析(5.2 - TreeMap, 红黑树的插入)

上一篇文章中咱们介绍了 Map，TreeMap 的接口和内部的数据结构实现：红黑树的概念。今天文章的主要内容是介绍红黑树的核心操做之一，插入操做的代码实现。

在开始本文以前请确认本身掌握了上一篇文章中说起的相关知识，即平衡二叉树，Color Flip，Left/Right Rotation 。

平衡二叉树的插入

在上一篇文章中介绍了平衡二叉树的概念，这是一种通过排序的数据结构，那么它的插入逻辑是怎么样的呢？让咱们对照 TreeMap 的代码看一下。

TreeMap 经过 put 方法向容器内添加元素，put 方法的开始以下：

public V put(K key, V value) {
    Entry<K,V> t = root;
    if (t == null) {
        compare(key, key); // type (and possibly null) check

        root = new Entry<>(key, value, null);
        size = 1;
        modCount++;
        return null;
    }
......

方法签名很容易理解，就是须要添加的 key 与 value，都是泛型。一开始的 if 判断当前红黑树的根节点是否为空，若是为空的话就将当前添加的数据做为根节点。

若是当前根节点不为空，说明已经有了红黑树的数据结构，则会执行插入的逻辑，让咱们继续往下看。

int cmp;
Entry<K,V> parent;
// split comparator and comparable paths
Comparator<? super K> cpr = comparator;
if (cpr != null) {
    do {
        parent = t;
        cmp = cpr.compare(key, t.key);
        if (cmp < 0)
            t = t.left;
        else if (cmp > 0)
            t = t.right;
        else
            return t.setValue(value);
    } while (t != null);
}

这部分主要是检查是否经过构造函数定义了 Comparator 对象，用以定义 key 的比较逻辑。若是你看下如下 if 对应的 else 分支，就会发现其实插入的逻辑都是相同的，所以咱们就分析 if 分支。

很容易看到插入的核心逻辑在 do...while 循环中，经过 compare 方法来比较 key 的大小。经过 parent 做为临时变量保存遍历树节点的当前节点的父节点。若是 key 小于当前节点的则向左遍历，不然向右，若是相等则直接将当前节点的值设为最新的 value。

当知足循环终止的条件，即 t == null 时，t 变量确定是 null ，而 parent 则指向 t 的父节点，此时程序已经找到在树中须要插入节点的位置。

接着就能够执行真正的插入操做，接着往下看。

Entry<K,V> e = new Entry<>(key, value, parent);
if (cmp < 0)
    parent.left = e;
else
    parent.right = e;
fixAfterInsertion(e);

很简单，经过 key 和 value 初始化 Entry 这是上一篇文章提到的红黑树节点的数据结构。而后按照比较结果的大小，设置插入的节点为左节点仍是右节点。到这里平衡二叉树插入的程序就结束了，很简单吧。而关键的是在 fixAfterInsertion 这个方法中，确认本身明白了插入逻辑后，再继续往下看。

红黑树插入后的调整

从上面的代码能够看出数据插入后有可能再也不符合平衡二叉树的定义，所以须要调整插入后节点的位置。同时 TreeMap 中使用的是红黑树结构，因此调整后的树状结构也应该符合红黑树的定义，而这部分的功能就是在 fixAfterInsertion 中。在阅读 fixAfterInsertion 的源码以前，先让咱们了解一下红黑树插入后调整的算法。

插入后调整

这部分的算法在 wikipedia 上的描述很简答，也易于理解，因此我引用 wikipedia 的描述来解释这部分算法。

为了以后描述方便，咱们定义几种节点类型和对应的缩写，具体以下：

• 当前节点：N
• 当前节点的父节点：P
• 当前节点的的兄弟节点，即当前节点父节点的另外一个子节点：S
• 当前节点父节点的兄弟节点，也称之为“叔叔”节点：U
• 当前节点的父节点的父节点，也称之为“祖父”节点：G

下面这幅图解释了各类节点类型的位置，请确认本身的理解正确，再接续。

首先按照插入节点的不一样状况进行对应的处理，具体分为如下 4 种状态：

1. N 为根节点
2. P 为黑色节点
3. P 为红色，而 U 也为红色
4. P 为红色，而 U 不为红色

这 4 种状态的处理其实很是简单，你最终发现其实只须要记住如何处理第 3，第 4 种状态就好了。让咱们开始吧。

第 1 种状态很是简单，做为根节点须要作的就是将当前节点的颜色变为黑色便可，其余什么都不用作。

第 2 种状态更简单，你什么都不用作 ^o^，只须要保证当前插入节点的颜色为红色便可。

第 3 种状态则须要作一些操做。还记得 Color Flip 操做吗？要作的第一步是对 G 节点作 Color Flip 操做，即将 G 节点变为红色，P 和 U 节点变为黑色。第二步是将 G 做为参数再次执行调整算法 ，能够看做是个递归调用。

第 4 种状态是最为复杂的一种，分为两个阶段，但整体来讲也很容易理解，放松心情往下看。

1. 对 **P** 进行 Rotation
    1. 若是 **N** 为 **P**  右节点，且 **P** 为 **G** 的左节点，则对 **P** 进行 Left Rotation
    2. 若是 **N** 为 **P** 的左节点，且 **P** 为 **G** 的右节点，则对 **P** 进行 Right Rotation
    3. 若是不知足上述的条件则什么都不作
2. 对 **G** 进行 Rotation 
    1. 若是 **N** 为 **P**  右节点，则对 **G** 进行 Left Rotation
    2. 若是 **N** 为 **P**  左节点，则对 **G** 进行 Right Rotation
    3. 将 **P** 的颜色变为**黑色**
    4. 将 **G** 变为**红色**

fixAfterInsertion 源码

先看一下 fixAfterInsertion 的源码：

/** From CLR */
private void fixAfterInsertion(Entry<K,V> x) {
    x.color = RED;

    while (x != null && x != root && x.parent.color == RED) {
        if (parentOf(x) == leftOf(parentOf(parentOf(x)))) {
            Entry<K,V> y = rightOf(parentOf(parentOf(x)));
            if (colorOf(y) == RED) {
                setColor(parentOf(x), BLACK);
                setColor(y, BLACK);
                setColor(parentOf(parentOf(x)), RED);
                x = parentOf(parentOf(x));
            } else {
                if (x == rightOf(parentOf(x))) {
                    x = parentOf(x);
                    rotateLeft(x);
                }
                setColor(parentOf(x), BLACK);
                setColor(parentOf(parentOf(x)), RED);
                rotateRight(parentOf(parentOf(x)));
            }
        } else {
            Entry<K,V> y = leftOf(parentOf(parentOf(x)));
            if (colorOf(y) == RED) {
                setColor(parentOf(x), BLACK);
                setColor(y, BLACK);
                setColor(parentOf(parentOf(x)), RED);
                x = parentOf(parentOf(x));
            } else {
                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;
}

这部分算法的实现来自于大名鼎鼎，大部分人半途而废的算法名著：<<算法导论>>，即注释中的 CLR(CLR 是算法导论 3 位做者名称的缩写)。

在开始阶段会将节点的颜色设置为红色，而后开始循环。循环条件很简单，当前节点为不为空，当前节点不为根节点，当前节点父节点的颜色为红色。其实分析一下这个条件就可看到已经覆盖了以前 4 种状况的第 1 ，第 2 种状况了。

接着看 if 分支的条件，if (parentOf(x) == leftOf(parentOf(parentOf(x))))，即 P 是 G 的左节点。紧接着的 Entry<K,V> y = rightOf(parentOf(parentOf(x))); 是获取 U 节点，即父节点的兄弟节点。而后判断 U 节点的颜色是否为红色。这时对照咱们提到的 4 种状态，你会发现此时程序的状态已经知足第 3 种状态了，即 P 为红色，U 也为红色。接着让咱们看看，在这种分支下的处理逻辑：

if (colorOf(y) == RED) {
    setColor(parentOf(x), BLACK);
    setColor(y, BLACK);
    setColor(parentOf(parentOf(x)), RED);
    x = parentOf(parentOf(x));
}

一开始的 3 行的三行就是 Color Flip 的操做，将 G 设为红色，P 与 U 设为黑色，而后将 x 指向 G，开始下一轮的循环。这彻底符合咱们提到第 3 种状况的算法逻辑。

接着看对应的 else 分支代码。

else {
    if (x == rightOf(parentOf(x))) {
        x = parentOf(x);
        rotateLeft(x);
    }
    setColor(parentOf(x), BLACK);
    setColor(parentOf(parentOf(x)), RED);
    rotateRight(parentOf(parentOf(x)));
}

这个分支走的就是咱们说起的第 4 种状态的分支了。一样的，此时 P 是 G 的左节点，而后 x == rightOf(parentOf(x)) 是检查 N 是否为 P 的右节点。知足该条件的话，按照第 4 种状态的第 1 阶段算法，应该将 P 节点作 Left Rotation，代码中也是如此作的。以后的三行代码：

setColor(parentOf(x), BLACK);
setColor(parentOf(parentOf(x)), RED);
rotateRight(parentOf(parentOf(x)));

就是第 4 种状态的第 2 阶段操做：设置 P 和 G 的颜色，并对 G 进行 Rotation。

对应最外层 if 分支的 else 分支的代码其实大部分是同样的，只是 rotation 旋转的方向不一样，对应第 4 种状况的另一个分支，我就不在解释，留给你本身从代码对应算法描述了。

小结

本次文章结合 TreeMap 的代码解释了红黑树插入和从新平衡的操做，你们能够认真的对照代码和算法描述理清思路，了解红黑树的数据结构特色和算法，下一篇文章会介绍红黑树的删除操做，这也是 TreeMap 和红黑树的最后一篇了，但愿这 3 篇文章可以让你真正掌握红黑树的算法。

欢迎关注个人微信号「且把金针度与人」,获取更多高质量文章