死磕 java集合之TreeMap源碼分析（二）- 內(nèi)含紅黑樹(shù)分析全過(guò)程

插入元素

插入元素，如果元素在樹(shù)中存在，則替換value；如果元素不存在，則插入到對(duì)應(yīng)的位置，再平衡樹(shù)。

public V put(K key, V value) {
    Entry<K,V> t = root;
    if (t == null) {
        // 如果沒(méi)有根節(jié)點(diǎn)，直接插入到根節(jié)點(diǎn)
        compare(key, key); // type (and possibly null) check
        root = new Entry<>(key, value, null);
        size = 1;
        modCount++;
        return null;
    }
    // key比較的結(jié)果
    int cmp;
    // 用來(lái)尋找待插入節(jié)點(diǎn)的父節(jié)點(diǎn)
    Entry<K,V> parent;
    // 根據(jù)是否有comparator使用不同的分支
    Comparator<? super K> cpr = comparator;
    if (cpr != null) {
        // 如果使用的是comparator方式，key值可以為null，只要在comparator.compare()中允許即可
        // 從根節(jié)點(diǎn)開(kāi)始遍歷尋找
        do {
            parent = t;
            cmp = cpr.compare(key, t.key);
            if (cmp < 0)
                // 如果小于0從左子樹(shù)尋找
                t = t.left;
            else if (cmp > 0)
                // 如果大于0從右子樹(shù)尋找
                t = t.right;
            else
                // 如果等于0，說(shuō)明插入的節(jié)點(diǎn)已經(jīng)存在了，直接更換其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;
        // 從根節(jié)點(diǎn)開(kāi)始遍歷尋找
        do {
            parent = t;
            cmp = k.compareTo(t.key);
            if (cmp < 0)
                // 如果小于0從左子樹(shù)尋找
                t = t.left;
            else if (cmp > 0)
                // 如果大于0從右子樹(shù)尋找
                t = t.right;
            else
                // 如果等于0，說(shuō)明插入的節(jié)點(diǎn)已經(jīng)存在了，直接更換其value值并返回舊值
                return t.setValue(value);
        } while (t != null);
    }
    // 如果沒(méi)找到，那么新建一個(gè)節(jié)點(diǎn)，并插入到樹(shù)中
    Entry<K,V> e = new Entry<>(key, value, parent);
    if (cmp < 0)
        // 如果小于0插入到左子節(jié)點(diǎn)
        parent.left = e;
    else
        // 如果大于0插入到右子節(jié)點(diǎn)
        parent.right = e;

    // 插入之后的平衡
    fixAfterInsertion(e);
    // 元素個(gè)數(shù)加1（不需要擴(kuò)容）
    size++;
    // 修改次數(shù)加1
    modCount++;
    // 如果插入了新節(jié)點(diǎn)返回空
    return null;
}

插入再平衡

插入的元素默認(rèn)都是紅色，因?yàn)椴迦爰t色元素只違背了第4條特性，那么我們只要根據(jù)這個(gè)特性來(lái)平衡就容易多了。

根據(jù)不同的情況有以下幾種處理方式：

1. 插入的元素如果是根節(jié)點(diǎn)，則直接涂成黑色即可，不用平衡；

2. 插入的元素的父節(jié)點(diǎn)如果為黑色，不需要平衡；

3. 插入的元素的父節(jié)點(diǎn)如果為紅色，則違背了特性4，需要平衡，平衡時(shí)又分成下面三種情況：

（如果父節(jié)點(diǎn)是祖父節(jié)點(diǎn)的左節(jié)點(diǎn)）

（如果父節(jié)點(diǎn)是祖父節(jié)點(diǎn)的右節(jié)點(diǎn)，則正好與上面反過(guò)來(lái)）

讓我們來(lái)看看TreeMap中的實(shí)現(xiàn)：

/**
 * 插入再平衡
 *（1）每個(gè)節(jié)點(diǎn)或者是黑色，或者是紅色。
 *（2）根節(jié)點(diǎn)是黑色。
 *（3）每個(gè)葉子節(jié)點(diǎn)（NIL）是黑色。（注意：這里葉子節(jié)點(diǎn)，是指為空(NIL或NULL)的葉子節(jié)點(diǎn)！）
 *（4）如果一個(gè)節(jié)點(diǎn)是紅色的，則它的子節(jié)點(diǎn)必須是黑色的。
 *（5）從一個(gè)節(jié)點(diǎn)到該節(jié)點(diǎn)的子孫節(jié)點(diǎn)的所有路徑上包含相同數(shù)目的黑節(jié)點(diǎn)。
 */
private void fixAfterInsertion(Entry<K,V> x) {
    // 插入的節(jié)點(diǎn)為紅節(jié)點(diǎn)，x為當(dāng)前節(jié)點(diǎn)
    x.color = RED;

    // 只有當(dāng)插入節(jié)點(diǎn)不是根節(jié)點(diǎn)且其父節(jié)點(diǎn)為紅色時(shí)才需要平衡（違背了特性4）
    while (x != null && x != root && x.parent.color == RED) {
        if (parentOf(x) == leftOf(parentOf(parentOf(x)))) {
            // a）如果父節(jié)點(diǎn)是祖父節(jié)點(diǎn)的左節(jié)點(diǎn)
            // y為叔叔節(jié)點(diǎn)
            Entry<K,V> y = rightOf(parentOf(parentOf(x)));
            if (colorOf(y) == RED) {
                // 情況1）如果叔叔節(jié)點(diǎn)為紅色
                // （1）將父節(jié)點(diǎn)設(shè)為黑色
                setColor(parentOf(x), BLACK);
                // （2）將叔叔節(jié)點(diǎn)設(shè)為黑色
                setColor(y, BLACK);
                // （3）將祖父節(jié)點(diǎn)設(shè)為紅色
                setColor(parentOf(parentOf(x)), RED);
                // （4）將祖父節(jié)點(diǎn)設(shè)為新的當(dāng)前節(jié)點(diǎn)
                x = parentOf(parentOf(x));
            } else {
                // 如果叔叔節(jié)點(diǎn)為黑色
                // 情況2）如果當(dāng)前節(jié)點(diǎn)為其父節(jié)點(diǎn)的右節(jié)點(diǎn)
                if (x == rightOf(parentOf(x))) {
                    // （1）將父節(jié)點(diǎn)設(shè)為當(dāng)前節(jié)點(diǎn)
                    x = parentOf(x);
                    // （2）以新當(dāng)前節(jié)點(diǎn)左旋
                    rotateLeft(x);
                }
                // 情況3）如果當(dāng)前節(jié)點(diǎn)為其父節(jié)點(diǎn)的左節(jié)點(diǎn)（如果是情況2）則左旋之后新當(dāng)前節(jié)點(diǎn)正好為其父節(jié)點(diǎn)的左節(jié)點(diǎn)了）
                // （1）將父節(jié)點(diǎn)設(shè)為黑色
                setColor(parentOf(x), BLACK);
                // （2）將祖父節(jié)點(diǎn)設(shè)為紅色
                setColor(parentOf(parentOf(x)), RED);
                // （3）以祖父節(jié)點(diǎn)為支點(diǎn)進(jìn)行右旋
                rotateRight(parentOf(parentOf(x)));
            }
        } else {
            // b）如果父節(jié)點(diǎn)是祖父節(jié)點(diǎn)的右節(jié)點(diǎn)
            // y是叔叔節(jié)點(diǎn)
            Entry<K,V> y = leftOf(parentOf(parentOf(x)));
            if (colorOf(y) == RED) {
                // 情況1）如果叔叔節(jié)點(diǎn)為紅色
                // （1）將父節(jié)點(diǎn)設(shè)為黑色
                setColor(parentOf(x), BLACK);
                // （2）將叔叔節(jié)點(diǎn)設(shè)為黑色
                setColor(y, BLACK);
                // （3）將祖父節(jié)點(diǎn)設(shè)為紅色
                setColor(parentOf(parentOf(x)), RED);
                // （4）將祖父節(jié)點(diǎn)設(shè)為新的當(dāng)前節(jié)點(diǎn)
                x = parentOf(parentOf(x));
            } else {
                // 如果叔叔節(jié)點(diǎn)為黑色
                // 情況2）如果當(dāng)前節(jié)點(diǎn)為其父節(jié)點(diǎn)的左節(jié)點(diǎn)
                if (x == leftOf(parentOf(x))) {
                    // （1）將父節(jié)點(diǎn)設(shè)為當(dāng)前節(jié)點(diǎn)
                    x = parentOf(x);
                    // （2）以新當(dāng)前節(jié)點(diǎn)右旋
                    rotateRight(x);
                }
                // 情況3）如果當(dāng)前節(jié)點(diǎn)為其父節(jié)點(diǎn)的右節(jié)點(diǎn)（如果是情況2）則右旋之后新當(dāng)前節(jié)點(diǎn)正好為其父節(jié)點(diǎn)的右節(jié)點(diǎn)了）
                // （1）將父節(jié)點(diǎn)設(shè)為黑色
                setColor(parentOf(x), BLACK);
                // （2）將祖父節(jié)點(diǎn)設(shè)為紅色
                setColor(parentOf(parentOf(x)), RED);
                // （3）以祖父節(jié)點(diǎn)為支點(diǎn)進(jìn)行左旋
                rotateLeft(parentOf(parentOf(x)));
            }
        }
    }
    // 平衡完成后將根節(jié)點(diǎn)設(shè)為黑色
    root.color = BLACK;
}

插入元素舉例

我們依次向紅黑樹(shù)中插入 4、2、3 三個(gè)元素，來(lái)一起看看整個(gè)紅黑樹(shù)平衡的過(guò)程。

三個(gè)元素都插入完成后，符合父節(jié)點(diǎn)是祖父節(jié)點(diǎn)的左節(jié)點(diǎn)，叔叔節(jié)點(diǎn)為黑色，且當(dāng)前節(jié)點(diǎn)是其父節(jié)點(diǎn)的右節(jié)點(diǎn)，即情況2）。

情況2）需要做以下兩步處理：

（1）將父節(jié)點(diǎn)作為新的當(dāng)前節(jié)點(diǎn)；

（2）以新當(dāng)節(jié)點(diǎn)為支點(diǎn)進(jìn)行左旋，進(jìn)入情況3）；

情況3）需要做以下三步處理：

（1）將父節(jié)點(diǎn)設(shè)為黑色；

（2）將祖父節(jié)點(diǎn)設(shè)為紅色；

（3）以祖父節(jié)點(diǎn)為支點(diǎn)進(jìn)行右旋，進(jìn)入下一次循環(huán)判斷；

下一次循環(huán)不符合父節(jié)點(diǎn)為紅色了，退出循環(huán)，插入再平衡完成。

未完待續(xù)，下一節(jié)我們一起探討紅黑樹(shù)刪除元素的操作。