Java Collections Framework source code analysis (5.2 - TreeMap, red-black tree insert)

The last article we introduced Map, TreeMapthe internal interfaces and data structures to achieve: the concept of red-black tree. Today the main contents of the article is to introduce one of the core operations of the red-black tree, insert the code.

Before starting this article, please confirm your master , " an article on " knowledge mentioned that balanced binary tree, Color Flip, Left / Right Rotation .

Balanced binary tree insertion

Introduces the concept of balanced binary tree in the previous article, this is a data structure to sort through, then its insertion logic is kind of how it? Let us control TreeMapcode to look at.

TreeMapBy putadding an element to the method of the container, putthe method begins as follows:

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;
    }
......
复制代码

The method signature is easy to understand, is the need to add key and value, are generic. A start ifdetermining whether the current red-black tree root node is empty, if the data is empty, the current will be added as the root node.

If the current root is not empty, explained that it had been red-black tree data structure, logic inserted will be executed, let's continue to look down.

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);
}
复制代码

This part is to check if the function is defined by constructing Comparatorthe object for a logical comparison of key definitions. If you look at the ifcorresponding elsebranch, you will find that in fact the inserted logic is the same, so we will analyze ifbranch.

很容易看到插入的核心逻辑在 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 篇文章能够让你真正掌握红黑树的算法。

本文使用 mdnice 排版

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