【死磕 Java 集合】--- TreeMap源碼分析（二）

插入元素

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

public
 V put(K key, V value) {

Entry
 t = root;

if
 (t == 
null
) {

// 如果沒有根節(jié)點，直接插入到根節(jié)點
        compare(key, key); 
// type (and possibly null) check
        root = 
new

Entry
<>(key, value, 
null
);
        size = 
1
;
        modCount++;

return

null
;
    }

// key比較的結(jié)果

int
 cmp;

// 用來尋找待插入節(jié)點的父節(jié)點

Entry
 parent;

// 根據(jù)是否有comparator使用不同的分支

Comparator
 cpr = comparator;

if
 (cpr != 
null
) {

// 如果使用的是comparator方式，key值可以為null，只要在comparator.compare()中允許即可

// 從根節(jié)點開始遍歷尋找

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，說明插入的節(jié)點已經(jīng)存在了，直接更換其value值并返回舊值

return
 t.setValue(value);
        } 
while
 (t != 
null
);
    }

else
 {

// 如果使用的是Comparable方式，key不能為null

if
 (key == 
null
)

throw

new

NullPointerException
();

@SuppressWarnings
(
"unchecked"
)

Comparable
 k = (
Comparable
) key;

// 從根節(jié)點開始遍歷尋找

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，說明插入的節(jié)點已經(jīng)存在了，直接更換其value值并返回舊值

return
 t.setValue(value);
        } 
while
 (t != 
null
);
    }

// 如果沒找到，那么新建一個節(jié)點，并插入到樹中

Entry
 e = 
new

Entry
<>(key, value, parent);

if
 (cmp < 
0
)

// 如果小于0插入到左子節(jié)點
        parent.left = e;

else

// 如果大于0插入到右子節(jié)點
        parent.right = e;

// 插入之后的平衡
    fixAfterInsertion(e);

// 元素個數(shù)加1（不需要擴容）
    size++;

// 修改次數(shù)加1
    modCount++;

// 如果插入了新節(jié)點返回空

return

null
;
}

插入再平衡

插入的元素默認(rèn)都是紅色，因為插入紅色元素只違背了第4條特性，那么我們只要根據(jù)這個特性來平衡就容易多了。

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

1. 
   
2. 插入的元素如果是根節(jié)點，則直接涂成黑色即可，不用平衡；
3. 
   
4. 插入的元素的父節(jié)點如果為黑色，不需要平衡；
5. 
   
6. 插入的元素的父節(jié)點如果為紅色，則違背了特性4，需要平衡，平衡時又分成下面三種情況：

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

情況策略1）父節(jié)點為紅色，叔叔節(jié)點也為紅色（1）將父節(jié)點設(shè)為黑色；

（2）將叔叔節(jié)點設(shè)為黑色；

（3）將祖父節(jié)點設(shè)為紅色；

（4）將祖父節(jié)點設(shè)為新的當(dāng)前節(jié)點，進(jìn)入下一次循環(huán)判斷；2）父節(jié)點為紅色，叔叔節(jié)點為黑色，且當(dāng)前節(jié)點是其父節(jié)點的右節(jié)點（1）將父節(jié)點作為新的當(dāng)前節(jié)點；

（2）以新當(dāng)節(jié)點為支點進(jìn)行左旋，進(jìn)入情況3）；3）父節(jié)點為紅色，叔叔節(jié)點為黑色，且當(dāng)前節(jié)點是其父節(jié)點的左節(jié)點（1）將父節(jié)點設(shè)為黑色；

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

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

（如果父節(jié)點是祖父節(jié)點的右節(jié)點，則正好與上面反過來）

情況策略1）父節(jié)點為紅色，叔叔節(jié)點也為紅色（1）將父節(jié)點設(shè)為黑色；

（2）將叔叔節(jié)點設(shè)為黑色；

（3）將祖父節(jié)點設(shè)為紅色；

（4）將祖父節(jié)點設(shè)為新的當(dāng)前節(jié)點，進(jìn)入下一次循環(huán)判斷；2）父節(jié)點為紅色，叔叔節(jié)點為黑色，且當(dāng)前節(jié)點是其父節(jié)點的左節(jié)點（1）將父節(jié)點作為新的當(dāng)前節(jié)點；

（2）以新當(dāng)節(jié)點為支點進(jìn)行右旋；3）父節(jié)點為紅色，叔叔節(jié)點為黑色，且當(dāng)前節(jié)點是其父節(jié)點的右節(jié)點（1）將父節(jié)點設(shè)為黑色；

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

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

讓我們來看看TreeMap中的實現(xiàn)：

/**
 * 插入再平衡
 *（1）每個節(jié)點或者是黑色，或者是紅色。
 *（2）根節(jié)點是黑色。
 *（3）每個葉子節(jié)點（NIL）是黑色。（注意：這里葉子節(jié)點，是指為空(NIL或NULL)的葉子節(jié)點！）
 *（4）如果一個節(jié)點是紅色的，則它的子節(jié)點必須是黑色的。
 *（5）從一個節(jié)點到該節(jié)點的子孫節(jié)點的所有路徑上包含相同數(shù)目的黑節(jié)點。
 */
private

void
 fixAfterInsertion(
Entry
 x) {

// 插入的節(jié)點為紅節(jié)點，x為當(dāng)前節(jié)點
    x.color = RED;

// 只有當(dāng)插入節(jié)點不是根節(jié)點且其父節(jié)點為紅色時才需要平衡（違背了特性4）

while
 (x != 
null
 && x != root && x.parent.color == RED) {

if
 (parentOf(x) == leftOf(parentOf(parentOf(x)))) {

// a）如果父節(jié)點是祖父節(jié)點的左節(jié)點

// y為叔叔節(jié)點

Entry
 y = rightOf(parentOf(parentOf(x)));

if
 (colorOf(y) == RED) {

// 情況1）如果叔叔節(jié)點為紅色

// （1）將父節(jié)點設(shè)為黑色
                setColor(parentOf(x), BLACK);

// （2）將叔叔節(jié)點設(shè)為黑色
                setColor(y, BLACK);

// （3）將祖父節(jié)點設(shè)為紅色
                setColor(parentOf(parentOf(x)), RED);

// （4）將祖父節(jié)點設(shè)為新的當(dāng)前節(jié)點
                x = parentOf(parentOf(x));
            } 
else
 {

// 如果叔叔節(jié)點為黑色

// 情況2）如果當(dāng)前節(jié)點為其父節(jié)點的右節(jié)點

if
 (x == rightOf(parentOf(x))) {

// （1）將父節(jié)點設(shè)為當(dāng)前節(jié)點
                    x = parentOf(x);

// （2）以新當(dāng)前節(jié)點左旋
                    rotateLeft(x);
                }

// 情況3）如果當(dāng)前節(jié)點為其父節(jié)點的左節(jié)點（如果是情況2）則左旋之后新當(dāng)前節(jié)點正好為其父節(jié)點的左節(jié)點了）

// （1）將父節(jié)點設(shè)為黑色
                setColor(parentOf(x), BLACK);

// （2）將祖父節(jié)點設(shè)為紅色
                setColor(parentOf(parentOf(x)), RED);

// （3）以祖父節(jié)點為支點進(jìn)行右旋
                rotateRight(parentOf(parentOf(x)));
            }
        } 
else
 {

// b）如果父節(jié)點是祖父節(jié)點的右節(jié)點

// y是叔叔節(jié)點

Entry
 y = leftOf(parentOf(parentOf(x)));

if
 (colorOf(y) == RED) {

// 情況1）如果叔叔節(jié)點為紅色

// （1）將父節(jié)點設(shè)為黑色
                setColor(parentOf(x), BLACK);

// （2）將叔叔節(jié)點設(shè)為黑色
                setColor(y, BLACK);

// （3）將祖父節(jié)點設(shè)為紅色
                setColor(parentOf(parentOf(x)), RED);

// （4）將祖父節(jié)點設(shè)為新的當(dāng)前節(jié)點
                x = parentOf(parentOf(x));
            } 
else
 {

// 如果叔叔節(jié)點為黑色

// 情況2）如果當(dāng)前節(jié)點為其父節(jié)點的左節(jié)點

if
 (x == leftOf(parentOf(x))) {

// （1）將父節(jié)點設(shè)為當(dāng)前節(jié)點
                    x = parentOf(x);

// （2）以新當(dāng)前節(jié)點右旋
                    rotateRight(x);
                }

// 情況3）如果當(dāng)前節(jié)點為其父節(jié)點的右節(jié)點（如果是情況2）則右旋之后新當(dāng)前節(jié)點正好為其父節(jié)點的右節(jié)點了）

// （1）將父節(jié)點設(shè)為黑色
                setColor(parentOf(x), BLACK);

// （2）將祖父節(jié)點設(shè)為紅色
                setColor(parentOf(parentOf(x)), RED);

// （3）以祖父節(jié)點為支點進(jìn)行左旋
                rotateLeft(parentOf(parentOf(x)));
            }
        }
    }

// 平衡完成后將根節(jié)點設(shè)為黑色
    root.color = BLACK;
}

插入元素舉例

我們依次向紅黑樹中插入 4、2、3 三個元素，來一起看看整個紅黑樹平衡的過程。

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

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

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

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

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

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

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

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

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

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