死磕 java集合之TreeMap源码分析（二）——红黑树全解析

欢迎关注个人公众号“彤哥读源码”，查看更多源码系列文章, 与彤哥一块儿畅游源码的海洋。

---

插入元素

插入元素，若是元素在树中存在，则替换value；若是元素不存在，则插入到对应的位置，再平衡树。

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比较的结果
    int cmp;
    // 用来寻找待插入节点的父节点
    Entry<K,V> parent;
    // 根据是否有comparator使用不一样的分支
    Comparator<? super K> cpr = comparator;
    if (cpr != null) {
        // 若是使用的是comparator方式，key值能够为null，只要在comparator.compare()中容许便可
        // 从根节点开始遍历寻找
        do {
            parent = t;
            cmp = cpr.compare(key, t.key);
            if (cmp < 0)
                // 若是小于0从左子树寻找
                t = t.left;
            else if (cmp > 0)
                // 若是大于0从右子树寻找
                t = t.right;
            else
                // 若是等于0，说明插入的节点已经存在了，直接更换其value值并返回旧值
                return t.setValue(value);
        } while (t != null);
    }
    else {
        // 若是使用的是Comparable方式，key不能为null
        if (key == null)
            throw new NullPointerException();
        @SuppressWarnings("unchecked")
        Comparable<? super K> k = (Comparable<? super K>) key;
        // 从根节点开始遍历寻找
        do {
            parent = t;
            cmp = k.compareTo(t.key);
            if (cmp < 0)
                // 若是小于0从左子树寻找
                t = t.left;
            else if (cmp > 0)
                // 若是大于0从右子树寻找
                t = t.right;
            else
                // 若是等于0，说明插入的节点已经存在了，直接更换其value值并返回旧值
                return t.setValue(value);
        } while (t != null);
    }
    // 若是没找到，那么新建一个节点，并插入到树中
    Entry<K,V> e = new Entry<>(key, value, parent);
    if (cmp < 0)
        // 若是小于0插入到左子节点
        parent.left = e;
    else
        // 若是大于0插入到右子节点
        parent.right = e;

    // 插入以后的平衡
    fixAfterInsertion(e);
    // 元素个数加1（不须要扩容）
    size++;
    // 修改次数加1
    modCount++;
    // 若是插入了新节点返回空
    return null;
}
复制代码

插入再平衡

插入的元素默认都是红色，由于插入红色元素只违背了第4条特性，那么咱们只要根据这个特性来平衡就容易多了。

根据不一样的状况有如下几种处理方式：

1. 插入的元素若是是根节点，则直接涂成黑色便可，不用平衡；

2. 插入的元素的父节点若是为黑色，不须要平衡；

3. 插入的元素的父节点若是为红色，则违背了特性4，须要平衡，平衡时又分红下面三种状况：

（若是父节点是祖父节点的左节点）

状况	策略
1）父节点为红色，叔叔节点也为红色	（1）将父节点设为黑色； （2）将叔叔节点设为黑色； （3）将祖父节点设为红色； （4）将祖父节点设为新的当前节点，进入下一次循环判断；
2）父节点为红色，叔叔节点为黑色，且当前节点是其父节点的右节点	（1）将父节点做为新的当前节点； （2）以新当节点为支点进行左旋，进入状况3）；
3）父节点为红色，叔叔节点为黑色，且当前节点是其父节点的左节点	（1）将父节点设为黑色； （2）将祖父节点设为红色； （3）以祖父节点为支点进行右旋，进入下一次循环判断；

（若是父节点是祖父节点的右节点，则正好与上面反过来）

状况	策略
1）父节点为红色，叔叔节点也为红色	（1）将父节点设为黑色； （2）将叔叔节点设为黑色； （3）将祖父节点设为红色； （4）将祖父节点设为新的当前节点，进入下一次循环判断；
2）父节点为红色，叔叔节点为黑色，且当前节点是其父节点的左节点	（1）将父节点做为新的当前节点； （2）以新当节点为支点进行右旋；
3）父节点为红色，叔叔节点为黑色，且当前节点是其父节点的右节点	（1）将父节点设为黑色； （2）将祖父节点设为红色； （3）以祖父节点为支点进行左旋，进入下一次循环判断；

让咱们来看看TreeMap中的实现：

/** * 插入再平衡 *（1）每一个节点或者是黑色，或者是红色。 *（2）根节点是黑色。 *（3）每一个叶子节点（NIL）是黑色。（注意：这里叶子节点，是指为空(NIL或NULL)的叶子节点！） *（4）若是一个节点是红色的，则它的子节点必须是黑色的。 *（5）从一个节点到该节点的子孙节点的全部路径上包含相同数目的黑节点。 */
private void fixAfterInsertion(Entry<K,V> x) {
    // 插入的节点为红节点，x为当前节点
    x.color = RED;

    // 只有当插入节点不是根节点且其父节点为红色时才须要平衡（违背了特性4）
    while (x != null && x != root && x.parent.color == RED) {
        if (parentOf(x) == leftOf(parentOf(parentOf(x)))) {
            // a）若是父节点是祖父节点的左节点
            // y为叔叔节点
            Entry<K,V> y = rightOf(parentOf(parentOf(x)));
            if (colorOf(y) == RED) {
                // 状况1）若是叔叔节点为红色
                // （1）将父节点设为黑色
                setColor(parentOf(x), BLACK);
                // （2）将叔叔节点设为黑色
                setColor(y, BLACK);
                // （3）将祖父节点设为红色
                setColor(parentOf(parentOf(x)), RED);
                // （4）将祖父节点设为新的当前节点
                x = parentOf(parentOf(x));
            } else {
                // 若是叔叔节点为黑色
                // 状况2）若是当前节点为其父节点的右节点
                if (x == rightOf(parentOf(x))) {
                    // （1）将父节点设为当前节点
                    x = parentOf(x);
                    // （2）以新当前节点左旋
                    rotateLeft(x);
                }
                // 状况3）若是当前节点为其父节点的左节点（若是是状况2）则左旋以后新当前节点正好为其父节点的左节点了）
                // （1）将父节点设为黑色
                setColor(parentOf(x), BLACK);
                // （2）将祖父节点设为红色
                setColor(parentOf(parentOf(x)), RED);
                // （3）以祖父节点为支点进行右旋
                rotateRight(parentOf(parentOf(x)));
            }
        } else {
            // b）若是父节点是祖父节点的右节点
            // y是叔叔节点
            Entry<K,V> y = leftOf(parentOf(parentOf(x)));
            if (colorOf(y) == RED) {
                // 状况1）若是叔叔节点为红色
                // （1）将父节点设为黑色
                setColor(parentOf(x), BLACK);
                // （2）将叔叔节点设为黑色
                setColor(y, BLACK);
                // （3）将祖父节点设为红色
                setColor(parentOf(parentOf(x)), RED);
                // （4）将祖父节点设为新的当前节点
                x = parentOf(parentOf(x));
            } else {
                // 若是叔叔节点为黑色
                // 状况2）若是当前节点为其父节点的左节点
                if (x == leftOf(parentOf(x))) {
                    // （1）将父节点设为当前节点
                    x = parentOf(x);
                    // （2）以新当前节点右旋
                    rotateRight(x);
                }
                // 状况3）若是当前节点为其父节点的右节点（若是是状况2）则右旋以后新当前节点正好为其父节点的右节点了）
                // （1）将父节点设为黑色
                setColor(parentOf(x), BLACK);
                // （2）将祖父节点设为红色
                setColor(parentOf(parentOf(x)), RED);
                // （3）以祖父节点为支点进行左旋
                rotateLeft(parentOf(parentOf(x)));
            }
        }
    }
    // 平衡完成后将根节点设为黑色
    root.color = BLACK;
}
复制代码

插入元素举例

咱们依次向红黑树中插入 四、二、3 三个元素，来一块儿看看整个红黑树平衡的过程。

三个元素都插入完成后，符合父节点是祖父节点的左节点，叔叔节点为黑色，且当前节点是其父节点的右节点，即状况2）。

状况2）须要作如下两步处理：

（1）将父节点做为新的当前节点；

（2）以新当节点为支点进行左旋，进入状况3）；

状况3）须要作如下三步处理：

（1）将父节点设为黑色；

（2）将祖父节点设为红色；

（3）以祖父节点为支点进行右旋，进入下一次循环判断；

下一次循环不符合父节点为红色了，退出循环，插入再平衡完成。

---

未完待续，下一节咱们一块儿探讨红黑树删除元素的操做。

如今公众号文章没办法留言了，若是有什么疑问或者建议请直接在公众号给我留言。

---

欢迎关注个人公众号“彤哥读源码”，查看更多源码系列文章, 与彤哥一块儿畅游源码的海洋。