死磕 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）以祖父節點爲支點進行右旋，進入下一次循環判斷；

下一次循環不符合父節點爲紅色了，退出循環，插入再平衡完成。

---

未完待續，下一節咱們一塊兒探討紅黑樹刪除元素的操做。

如今公衆號文章沒辦法留言了，若是有什麼疑問或者建議請直接在公衆號給我留言。

---

歡迎關注個人公衆號「彤哥讀源碼」，查看更多源碼系列文章, 與彤哥一塊兒暢遊源碼的海洋。