死磕 java集合之TreeMap原始碼分析(三)——紅黑樹全解析

彤哥讀原始碼發表於2019-04-13

歡迎關注我的公眾號“彤哥讀原始碼”,檢視更多原始碼系列文章, 與彤哥一起暢遊原始碼的海洋。


刪除元素

刪除元素本身比較簡單,就是採用二叉樹的刪除規則。

(1)如果刪除的位置有兩個葉子節點,則從其右子樹中取最小的元素放到刪除的位置,然後把刪除位置移到替代元素的位置,進入下一步。

(2)如果刪除的位置只有一個葉子節點(有可能是經過第一步轉換後的刪除位置),則把那個葉子節點作為替代元素,放到刪除的位置,然後把這個葉子節點刪除。

(3)如果刪除的位置沒有葉子節點,則直接把這個刪除位置的元素刪除即可。

(4)針對紅黑樹,如果刪除位置是黑色節點,還需要做再平衡。

(5)如果有替代元素,則以替代元素作為當前節點進入再平衡。

(6)如果沒有替代元素,則以刪除的位置的元素作為當前節點進入再平衡,平衡之後再刪除這個節點。

public V remove(Object key) {
    // 獲取節點
    Entry<K,V> p = getEntry(key);
    if (p == null)
        return null;

    V oldValue = p.value;
    // 刪除節點
    deleteEntry(p);
    // 返回刪除的value
    return oldValue;
}

private void deleteEntry(Entry<K,V> p) {
    // 修改次數加1
    modCount++;
    // 元素個數減1
    size--;

    if (p.left != null && p.right != null) {
        // 如果當前節點既有左子節點,又有右子節點
        // 取其右子樹中最小的節點
        Entry<K,V> s = successor(p);
        // 用右子樹中最小節點的值替換當前節點的值
        p.key = s.key;
        p.value = s.value;
        // 把右子樹中最小節點設為當前節點
        p = s;
        // 這種情況實際上並沒有刪除p節點,而是把p節點的值改了,實際刪除的是p的後繼節點
    }

    // 如果原來的當前節點(p)有2個子節點,則當前節點已經變成原來p的右子樹中的最小節點了,也就是說其沒有左子節點了
    // 到這一步,p肯定只有一個子節點了
    // 如果當前節點有子節點,則用子節點替換當前節點
    Entry<K,V> replacement = (p.left != null ? p.left : p.right);

    if (replacement != null) {
        // 把替換節點直接放到當前節點的位置上(相當於刪除了p,並把替換節點移動過來了)
        replacement.parent = p.parent;
        if (p.parent == null)
            root = replacement;
        else if (p == p.parent.left)
            p.parent.left  = replacement;
        else
            p.parent.right = replacement;

        // 將p的各項屬性都設為空
        p.left = p.right = p.parent = null;

        // 如果p是黑節點,則需要再平衡
        if (p.color == BLACK)
            fixAfterDeletion(replacement);
    } else if (p.parent == null) {
        // 如果當前節點就是根節點,則直接將根節點設為空即可
        root = null;
    } else {
        // 如果當前節點沒有子節點且其為黑節點,則把自己當作虛擬的替換節點進行再平衡
        if (p.color == BLACK)
            fixAfterDeletion(p);

        // 平衡完成後刪除當前節點(與父節點斷絕關係)
        if (p.parent != null) {
            if (p == p.parent.left)
                p.parent.left = null;
            else if (p == p.parent.right)
                p.parent.right = null;
            p.parent = null;
        }
    }
}
複製程式碼

刪除再平衡

經過上面的處理,真正刪除的肯定是黑色節點才會進入到再平衡階段。

因為刪除的是黑色節點,導致整顆樹不平衡了,所以這裡我們假設把刪除的黑色賦予當前節點,這樣當前節點除了它自已的顏色還多了一個黑色,那麼:

(1)如果當前節點是根節點,則直接塗黑即可,不需要再平衡;

(2)如果當前節點是紅+黑節點,則直接塗黑即可,不需要平衡;

(3)如果當前節點是黑+黑節點,則我們只要通過旋轉把這個多出來的黑色不斷的向上傳遞到一個紅色節點即可,這又可能會出現以下四種情況:

(假設當前節點為父節點的左子節點)

情況 策略
1)x是黑+黑節點,x的兄弟是紅節點 (1)將兄弟節點設為黑色;
(2)將父節點設為紅色;
(3)以父節點為支點進行左旋;
(4)重新設定x的兄弟節點,進入下一步;
2)x是黑+黑節點,x的兄弟是黑節點,且兄弟節點的兩個子節點都是黑色 (1)將兄弟節點設定為紅色;
(2)將x的父節點作為新的當前節點,進入下一次迴圈;
3)x是黑+黑節點,x的兄弟是黑節點,且兄弟節點的右子節點為黑色,左子節點為紅色 (1)將兄弟節點的左子節點設為黑色;
(2)將兄弟節點設為紅色;
(3)以兄弟節點為支點進行右旋;
(4)重新設定x的兄弟節點,進入下一步;
3)x是黑+黑節點,x的兄弟是黑節點,且兄弟節點的右子節點為紅色,左子節點任意顏色 (1)將兄弟節點的顏色設為父節點的顏色;
(2)將父節點設為黑色;
(3)將兄弟節點的右子節點設為黑色;
(4)以父節點為支點進行左旋;
(5)將root作為新的當前節點(退出迴圈);

(假設當前節點為父節點的右子節點,正好反過來)

情況 策略
1)x是黑+黑節點,x的兄弟是紅節點 (1)將兄弟節點設為黑色;
(2)將父節點設為紅色;
(3)以父節點為支點進行右旋;
(4)重新設定x的兄弟節點,進入下一步;
2)x是黑+黑節點,x的兄弟是黑節點,且兄弟節點的兩個子節點都是黑色 (1)將兄弟節點設定為紅色;
(2)將x的父節點作為新的當前節點,進入下一次迴圈;
3)x是黑+黑節點,x的兄弟是黑節點,且兄弟節點的左子節點為黑色,右子節點為紅色 (1)將兄弟節點的右子節點設為黑色;
(2)將兄弟節點設為紅色;
(3)以兄弟節點為支點進行左旋;
(4)重新設定x的兄弟節點,進入下一步;
3)x是黑+黑節點,x的兄弟是黑節點,且兄弟節點的左子節點為紅色,右子節點任意顏色 (1)將兄弟節點的顏色設為父節點的顏色;
(2)將父節點設為黑色;
(3)將兄弟節點的左子節點設為黑色;
(4)以父節點為支點進行右旋;
(5)將root作為新的當前節點(退出迴圈);

讓我們來看看TreeMap中的實現:

/**
 * 刪除再平衡
 *(1)每個節點或者是黑色,或者是紅色。
 *(2)根節點是黑色。
 *(3)每個葉子節點(NIL)是黑色。(注意:這裡葉子節點,是指為空(NIL或NULL)的葉子節點!)
 *(4)如果一個節點是紅色的,則它的子節點必須是黑色的。
 *(5)從一個節點到該節點的子孫節點的所有路徑上包含相同數目的黑節點。
 */
private void fixAfterDeletion(Entry<K,V> x) {
    // 只有當前節點不是根節點且當前節點是黑色時才進入迴圈
    while (x != root && colorOf(x) == BLACK) {
        if (x == leftOf(parentOf(x))) {
            // 如果當前節點是其父節點的左子節點
            // sib是當前節點的兄弟節點
            Entry<K,V> sib = rightOf(parentOf(x));

            // 情況1)如果兄弟節點是紅色
            if (colorOf(sib) == RED) {
                // (1)將兄弟節點設為黑色
                setColor(sib, BLACK);
                // (2)將父節點設為紅色
                setColor(parentOf(x), RED);
                // (3)以父節點為支點進行左旋
                rotateLeft(parentOf(x));
                // (4)重新設定x的兄弟節點,進入下一步
                sib = rightOf(parentOf(x));
            }

            if (colorOf(leftOf(sib))  == BLACK &&
                    colorOf(rightOf(sib)) == BLACK) {
                // 情況2)如果兄弟節點的兩個子節點都是黑色
                // (1)將兄弟節點設定為紅色
                setColor(sib, RED);
                // (2)將x的父節點作為新的當前節點,進入下一次迴圈
                x = parentOf(x);
            } else {
                if (colorOf(rightOf(sib)) == BLACK) {
                    // 情況3)如果兄弟節點的右子節點為黑色
                    // (1)將兄弟節點的左子節點設為黑色
                    setColor(leftOf(sib), BLACK);
                    // (2)將兄弟節點設為紅色
                    setColor(sib, RED);
                    // (3)以兄弟節點為支點進行右旋
                    rotateRight(sib);
                    // (4)重新設定x的兄弟節點
                    sib = rightOf(parentOf(x));
                }
                // 情況4)
                // (1)將兄弟節點的顏色設為父節點的顏色
                setColor(sib, colorOf(parentOf(x)));
                // (2)將父節點設為黑色
                setColor(parentOf(x), BLACK);
                // (3)將兄弟節點的右子節點設為黑色
                setColor(rightOf(sib), BLACK);
                // (4)以父節點為支點進行左旋
                rotateLeft(parentOf(x));
                // (5)將root作為新的當前節點(退出迴圈)
                x = root;
            }
        } else { // symmetric
            // 如果當前節點是其父節點的右子節點
            // sib是當前節點的兄弟節點
            Entry<K,V> sib = leftOf(parentOf(x));

            // 情況1)如果兄弟節點是紅色
            if (colorOf(sib) == RED) {
                // (1)將兄弟節點設為黑色
                setColor(sib, BLACK);
                // (2)將父節點設為紅色
                setColor(parentOf(x), RED);
                // (3)以父節點為支點進行右旋
                rotateRight(parentOf(x));
                // (4)重新設定x的兄弟節點
                sib = leftOf(parentOf(x));
            }

            if (colorOf(rightOf(sib)) == BLACK &&
                    colorOf(leftOf(sib)) == BLACK) {
                // 情況2)如果兄弟節點的兩個子節點都是黑色
                // (1)將兄弟節點設定為紅色
                setColor(sib, RED);
                // (2)將x的父節點作為新的當前節點,進入下一次迴圈
                x = parentOf(x);
            } else {
                if (colorOf(leftOf(sib)) == BLACK) {
                    // 情況3)如果兄弟節點的左子節點為黑色
                    // (1)將兄弟節點的右子節點設為黑色
                    setColor(rightOf(sib), BLACK);
                    // (2)將兄弟節點設為紅色
                    setColor(sib, RED);
                    // (3)以兄弟節點為支點進行左旋
                    rotateLeft(sib);
                    // (4)重新設定x的兄弟節點
                    sib = leftOf(parentOf(x));
                }
                // 情況4)
                // (1)將兄弟節點的顏色設為父節點的顏色
                setColor(sib, colorOf(parentOf(x)));
                // (2)將父節點設為黑色
                setColor(parentOf(x), BLACK);
                // (3)將兄弟節點的左子節點設為黑色
                setColor(leftOf(sib), BLACK);
                // (4)以父節點為支點進行右旋
                rotateRight(parentOf(x));
                // (5)將root作為新的當前節點(退出迴圈)
                x = root;
            }
        }
    }

    // 退出條件為多出來的黑色向上傳遞到了根節點或者紅節點
    // 則將x設為黑色即可滿足紅黑樹規則
    setColor(x, BLACK);
}
複製程式碼

刪除元素舉例

假設我們有下面這樣一顆紅黑樹。

treemap-delete1

我們刪除6號元素,則從右子樹中找到了最小元素7,7又沒有子節點了,所以把7作為當前節點進行再平衡。

我們看到7是黑節點,且其兄弟為黑節點,且其兄弟的兩個子節點都是紅色,滿足情況4),平衡之後如下圖所示。

treemap-delete2

我們再刪除7號元素,則從右子樹中找到了最小元素8,8有子節點且為黑色,所以8的子節點9是替代節點,以9為當前節點進行再平衡。

我們發現9是紅節點,則直接把它塗成黑色即滿足了紅黑樹的特性,不需要再過多的平衡了。

treemap-delete3

這次我們來個狠的,把根節點刪除,從右子樹中找到了最小的元素5,5沒有子節點,所以把5作為當前節點進行再平衡。

我們看到5是黑節點,且其兄弟為紅色,符合情況1),平衡之後如下圖所示,然後進入情況2)。

treemap-delete4

對情況2)進行再平衡後如下圖所示。

treemap-delete5

然後進入下一次迴圈,發現不符合迴圈條件了,直接把x塗為黑色即可,退出這個方法之後會把舊x刪除掉(見deleteEntry()方法),最後的結果就是下面這樣。

treemap-delete6


未完待續,下一節我們一起探討紅黑樹遍歷元素的操作。

現在公眾號文章沒辦法留言了,如果有什麼疑問或者建議請直接在公眾號給我留言。


歡迎關注我的公眾號“彤哥讀原始碼”,檢視更多原始碼系列文章, 與彤哥一起暢遊原始碼的海洋。

qrcode

相關文章