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 篇文章可以讓你真正掌握紅黑樹的算法。

歡迎關注個人微信號「且把金針度與人」,獲取更多高質量文章