TreeMap源碼詳解—徹底搞懂紅黑樹的平衡操作

介紹

TreeSet 和 TreeMap 在Java里有著相同的實現(xiàn)，前者僅僅是對后者做了一層包裝，也就是說 TreeSet里面有一個TreeMap (適配器模式)。

Java TreeMap 實現(xiàn)了 SortedMap 接口，也就是說會按照key的大小順序對 Map 中的元素進行排序，key大小的評判可以通過其本身的自然順序(natural ordering)，也可以通過構造時傳入的比較器(Comparator)。

TreeMap底層通過紅黑樹(Red-Black tree)實現(xiàn) ，也就意味著containsKey(), get(), put(), remove()都有著log(n)的時間復雜度。

1. TreeMap存儲 Key-Value 對時，需要根據(jù) key-value 對進行排序。TreeMap 可以保證所有的 Key-Value 對處于 有序狀態(tài)。

   • 正常情況下TreeMap是不能存入值為null的鍵的。

   • 但通過自定義比較器能讓TreeMap存入一個值為null的鍵。

   • 存入的值為null鍵對應的值不能通過通過它來獲取，只能通過直接遍歷Values。

2. TreeSet底層使用 紅黑樹結構存儲數(shù)據(jù)

3. TreeMap 的 Key 的排序：

   • 自然排序：TreeMap 的所有的 Key 必須實現(xiàn) Comparable 接口，而且所有的 Key 應該是同一個類的對象，否則將會拋出 ClasssCastException

   • 定制排序：創(chuàng)建 TreeMap 時，傳入一個 Comparator 對象，該對象負責對TreeMap 中的所有 key 進行排序。此時不需要 Map 的 Key 實現(xiàn)Comparable 接口

4. TreeMap判斷 兩個key 相等的標準：兩個key通過compareTo()方法或者compare()方法返回0。

底層實現(xiàn)

繼承接口的關鍵方法

public class TreeMap
    extends AbstractMap
    implements NavigableMap, Cloneable, java.io.Serializable

返回用于排序此映射中的鍵的比較器，如果此映射使用其鍵的自然排序，則返回null。

SortedMap接口：

Comparator comparator()
返回用于排序此映射中的鍵的比較器，如果此映射使用其鍵的自然排序，則返回null。

Set> entrySet()
返回此映射中包含的映射的Set視圖。

K firstKey()
返回當前映射中的第一個(最低)鍵。

K lastKey()
返回當前映射中的最后(最高)鍵。

NavigableMap接口：

Map.Entry ceilingEntry(K key)
返回與大于或等于給定鍵的最小鍵相關聯(lián)的鍵值映射，如果沒有這樣的鍵則返回null。

K ceilingKey(K key)
返回大于或等于給定鍵的最小鍵，如果沒有這樣的鍵，則返回null。

NavigableMap descendingMap()
返回此映射中包含的映射的倒序視圖。

Map.Entry firstEntry()
返回與該映射中最小的鍵關聯(lián)的鍵值映射，如果映射為空，則返回null。

Map.Entry floorEntry(K key)
返回與小于或等于給定鍵的最大鍵相關聯(lián)的鍵值映射，如果沒有這樣的鍵則返回null。

SortedMap headMap(K toKey)
返回該映射中鍵嚴格小于toKey的部分的視圖。

Map.Entry higherEntry(K key)
返回與嚴格大于給定鍵的最小鍵關聯(lián)的鍵值映射，如果沒有這樣的鍵，則返回null。

Map.Entry lastEntry()
返回與此映射中最大鍵關聯(lián)的鍵值映射，如果映射為空，則返回null。

Map.Entry lowerEntry(K key)
返回與嚴格小于給定鍵的最大鍵關聯(lián)的鍵值映射，如果沒有這樣的鍵，則返回null。

Map.Entry pollFirstEntry()
刪除并返回與該映射中最小的鍵關聯(lián)的鍵值映射，如果映射為空，則返回null。

Map.Entry pollLastEntry()
刪除并返回與此映射中最大鍵關聯(lián)的鍵值映射，如果映射為空，則返回null。

SortedMap subMap(K fromKey, K toKey)
返回該映射中鍵范圍從fromKey(包含)到toKey(獨占)的部分的視圖。

SortedMap tailMap(K fromKey)
返回該映射中鍵大于或等于fromKey的部分的視圖。

初始屬性

//比較器
private final Comparator comparator;

//root根的值
private transient Entry root;

//map中數(shù)據(jù)量
private transient int size = 0;

//修改次數(shù)
private transient int modCount = 0;

構造方法

public TreeMap() {
    comparator = null;
}

//有參構造，可以重新定義比較器
public TreeMap(Comparator comparator) {
    this.comparator = comparator;
}

//有參構造，將其他map用TreeMap存儲
public TreeMap(Map m) {
    comparator = null;
    putAll(m);
}

put 方法

public V put(K key, V value) {
    //將root賦值給局部變量
    Entry t = root;
    if (t == null) {//初始操作
        //檢查key是否為空
        compare(key, key); // type (and possibly null) check
        //將要添加的key value封裝為一個entry對象，并賦值給root
        root = new Entry<>(key, value, null);
        size = 1;
        modCount++;
        return null;
    }
    int cmp;
    Entry parent;//父節(jié)點
    // split comparator and comparable paths
    Comparator cpr = comparator;//獲取比較器
    if (cpr != null) {
        do {
            parent = t;//將root賦值給了parent
            cmp = cpr.compare(key, t.key);//和root節(jié)點比較大小
            if (cmp < 0)//key比根節(jié)點更小，t就使其為左節(jié)點
                t = t.left;
            else if (cmp > 0)//否則使其為右節(jié)點
                t = t.right;
            else
                return t.setValue(value);//大小相等，直接修改值
        } while (t != null);
    }
    else {
        if (key == null)
            throw new NullPointerException();
        @SuppressWarnings("unchecked")
            Comparable k = (Comparable) key;
        do {
            parent = t;
            cmp = k.compareTo(t.key);
            if (cmp < 0)
                t = t.left;
            else if (cmp > 0)
                t = t.right;
            else
                return t.setValue(value);
        } while (t != null);
    }
    // 到此 t 就是要插入節(jié)點的父節(jié)點，即parent
    //將 k v 對封裝成entry對象
    Entry e = new Entry<>(key, value, parent);
    if (cmp < 0)
        parent.left = e;//插入節(jié)點在 父節(jié)點 的左側
    else 
        parent.right = e;//插入節(jié)點在 父節(jié)點 的右側
    fixAfterInsertion(e);//實現(xiàn)紅黑樹的平衡
    size++;
    modCount++;
    return null;
}

關于紅黑樹的平衡，將在后文介紹

紅黑樹

紅黑樹的性質

• 每個節(jié)點要么是黑色，要么是紅色。

• 根節(jié)點是黑色。

• 每個葉子節(jié)點（NIL）是黑色。

• 每個紅色結點的兩個子結點一定都是黑色。

• 任意一結點(包含本身)到其葉子結點的路徑都包含數(shù)量相同的黑結點。

紅黑樹的優(yōu)點：

由于在AVL樹中，由于AVL樹是絕對平衡的，所有在進行插入和刪除的時候，為了維護其絕對的平衡性，有時候進行修改節(jié)點的操作，需要進行到根節(jié)點，旋轉的次數(shù)比較多，所以出現(xiàn)了紅黑樹，當數(shù)據(jù)不是靜態(tài)的數(shù)據(jù)而是動態(tài)的數(shù)據(jù)，進行插入和刪除的時候就不用去維護絕對的平衡，也就減少了旋轉的次數(shù)，照樣可以提高效率，并且紅黑樹的平均查找效率還是logn

紅黑樹平衡源碼

插入紅黑樹初始的顏色肯定為紅色，注意: 插入節(jié)點必須為紅色 ，理由很簡單，紅色在父節(jié)點(如果存在)為黑色節(jié)點時，紅黑樹的黑色平衡沒被破壞，不需要做自平衡操作。但如果插入結點是黑色，那么插入位置所在的子樹黑色結點總是多1，必須做自平衡。

可以先看TreeMap紅黑樹平衡源碼，也可以先看后文情景

private void fixAfterInsertion(Entry x) {
    x.color = RED;//插入時先把顏色設置為紅色
    //循環(huán)條件是x不等于空，x不是根，且父節(jié)點為紅色
    //原因在于①x為根的話可以直接置為黑色(對應情景1)；②父節(jié)點若為黑色，x直接以紅色插入，不影響平衡條件(對應情景3)
    while (x != null && x != root && x.parent.color == RED) {
        //判斷父節(jié)點  是否是  祖父節(jié)點的左側節(jié)點
        if (parentOf(x) == leftOf(parentOf(parentOf(x)))) {
            //獲取 祖父節(jié)點的右側節(jié)點，也就是叔叔節(jié)點S
            Entry y = rightOf(parentOf(parentOf(x)));
            //如果 叔叔節(jié)點為紅色(對應情景4.1) 
            if (colorOf(y) == RED) {
                //1.將P和S設置為黑色2.將PP設置為紅色3.將PP設置為當前插入節(jié)點再執(zhí)行規(guī)則
                setColor(parentOf(x), BLACK);
                setColor(y, BLACK);
                setColor(parentOf(parentOf(x)), RED);
                x = parentOf(parentOf(x));
            } else {//叔叔節(jié)點不存在(對應情景4.2)
              //如果x 是父節(jié)點的右節(jié)點(對應情景4.2.2)
                if (x == rightOf(parentOf(x))) {
                    x = parentOf(x);//將x的父節(jié)點P作為插入節(jié)點
                    rotateLeft(x);//左旋
                } //左旋完之后 插入節(jié)點x就是P的左節(jié)點
                //如果x是父節(jié)點的左節(jié)點；那就只做一次右旋(對應情景4.2.1)
                setColor(parentOf(x), BLACK);//將P設置為黑色
                setColor(parentOf(parentOf(x)), RED);//將PP設置為紅色
                rotateRight(parentOf(parentOf(x)));//右旋
            }
        } else {//父節(jié)點 是  祖父節(jié)點的右側節(jié)點
            //獲取 祖父節(jié)點的左側節(jié)點，也就是叔叔節(jié)點S
            Entry y = leftOf(parentOf(parentOf(x)));
            if (colorOf(y) == RED) {//(對應情景4.1) 
                setColor(parentOf(x), BLACK);
                setColor(y, BLACK);
                setColor(parentOf(parentOf(x)), RED);
                x = parentOf(parentOf(x));
            } else {//(對應情景4.3)
                if (x == leftOf(parentOf(x))) {//(對應情景4.3.2)
                    x = parentOf(x);
                    rotateRight(x);
                }
                //(對應情景4.3.1)
                setColor(parentOf(x), BLACK);
                setColor(parentOf(parentOf(x)), RED);
                rotateLeft(parentOf(parentOf(x)));
            }
        }
    }
    //根節(jié)點的顏色為黑色
    root.color = BLACK;
}

左旋：以某個節(jié)點作為支點(旋轉節(jié)點)，其右子節(jié)點變?yōu)樾�轉節(jié)點的父節(jié)點，右子節(jié)點的左子節(jié)點變?yōu)樾�轉節(jié)點的右子節(jié)點，旋轉節(jié)點的左子節(jié)點保持不變。右子節(jié)點的左子節(jié)點相當于從右子節(jié)點上“斷開”，重新連接到旋轉節(jié)點上。

//左旋
private void rotateLeft(Entry p) {
    if (p != null) {
        //
        Entry r = p.right;
        p.right = r.left;
        if (r.left != null)
            r.left.parent = p;
        r.parent = p.parent;
        if (p.parent == null)
            root = r;
        else if (p.parent.left == p)
            p.parent.left = r;
        else
            p.parent.right = r;
        r.left = p;
        p.parent = r;
    }
}

右旋：以某個節(jié)點作為支點(旋轉節(jié)點)，其左子節(jié)點變?yōu)樾�轉節(jié)點的父節(jié)點，左子節(jié)點的右子節(jié)點變?yōu)樾�轉節(jié)點的左子節(jié)點，旋轉節(jié)點的右子節(jié)點保持不變。左子節(jié)點的右子節(jié)點相當于從左子節(jié)點上“斷開”，重新連接到旋轉節(jié)點上。

private void rotateRight(Entry p) {
    if (p != null) {
        Entry l = p.left;
        p.left = l.right;
        if (l.right != null) l.right.parent = p;
        l.parent = p.parent;
        if (p.parent == null)
            root = l;
        else if (p.parent.right == p)
            p.parent.right = l;
        else p.parent.left = l;
        l.right = p;
        p.parent = l;
    }
}

以下情景中插入的操作轉載于： https://www.jianshu.com/p/e136ec79235c

情景1：紅黑樹為空樹

最簡單的一種情景，直接把插入結點作為根結點就行，但注意，根據(jù)紅黑樹性質2：根節(jié)點是黑色。還需要把插入結點設為黑色。

處理：把插入結點作為根結點，并把結點設置為黑色。

情景2：插入結點的Key已存在

插入結點的Key已存在，既然紅黑樹總保持平衡，在插入前紅黑樹已經(jīng)是平衡的，那么把插入結點設置為將要替代結點的顏色，再把結點的值更新就完成插入（這里的更新的其實是相同 key的 value）

處理：

• 把I設為當前結點的顏色

• 更新當前結點的值為插入結點的值

情景3：插入結點的父結點為黑結點

由于插入的結點是紅色的，當插入結點的黑色時，并不會影響紅黑樹的平衡，直接插入即可，無需做自平衡。

處理：直接插入。

情景4：插入結點的父結點為紅結點

再次回想下紅黑樹的性質2：根結點是黑色。如果插入的父結點為紅結點，那么該父結點不可能為根結點，所以插入結點總是存在祖父結點。這點很重要，因為后續(xù)的旋轉操作肯定需要祖父結點的參與。

情景4.1：叔叔結點存在并且為紅結點

從紅黑樹性質4可以，祖父結點肯定為黑結點，因為不可以同時存在兩個相連的紅結點。那么此時該插入子樹的紅黑層數(shù)的情況是：黑紅紅。顯然最簡單的處理方式是把其改為：紅黑紅。

(以下描述中I為插入節(jié)點，P為父節(jié)點，PP為祖父節(jié)點，S為叔叔節(jié)點)

處理：

• 將P和S設置為黑色

• 將PP設置為紅色

• 把PP設置為當前插入結點

把PP結點設為紅色了，如果PP的父結點是黑色，那么無需再做任何處理；但如果PP的父結點是紅色，根據(jù)性質4(每個紅色結點的兩個子結點一定都是黑色。)，此時紅黑樹已不平衡了，所以還需要把PP當作新的插入結點，繼續(xù)做插入操作自平衡處理，直到平衡為止。

試想下PP剛好為根結點時，那么根據(jù)性質2，必須把PP重新設為黑色，那么樹的紅黑結構變?yōu)椋汉诤诩t。換句話說，從根結點到葉子結點的路徑中，黑色結點增加了。這也是唯一一種會增加紅黑樹黑色結點層數(shù)的插入情景。

我們還可以總結出另外一個經(jīng)驗：紅黑樹的生長是自底向上的。這點不同于普通的二叉查找樹，普通的二叉查找樹的生長是自頂向下的。

情景4.2：叔叔結點不存在，并且插入結點的父親結點是祖父結點的左子結點

單純從插入前來看，也即不算情景4.1自底向上處理時的情況，叔叔結點非紅即為葉子結點(Nil)。

我們沒有考慮叔叔節(jié)點是黑節(jié)點情況，因為如果叔叔結點為黑結點，而父結點為紅結點，那么叔叔結點所在的子樹的黑色結點就比父結點所在子樹的多了，這不滿足紅黑樹的性質5。后續(xù)情景同樣如此，不再多做說明了。

情景4.2.1：插入結點是其父結點的左子結點

處理：

• 將P設為黑色

• 將PP設為紅色

• 對PP進行右旋

左邊兩個紅結點，右邊不存在，那么一邊一個剛剛好，并且因為為紅色，肯定不會破壞樹的平衡。

情景4.2.2：插入結點是其父結點的右子結點

這種情景顯然可以轉換為情景4.2.1

處理：

• 對P進行左旋

• 把P設置為插入結點，得到情景4.2.1

• 進行情景4.2.1的處理

情景4.3：叔叔結點不存在，并且插入結點的父親結點是祖父結點的右子結點

該情景對應情景4.2，只是方向反轉

情景4.3.1：插入結點是其父結點的右子結點

處理：

• 將P設為黑色

• 將PP設為紅色

• 對PP進行左旋

情景4.3.2：插入結點是其父結點的左子結點

處理：

• 對P進行右旋

• 把P設置為插入結點，得到情景4.3.1

• 進行情景4.3.1的處理

關于作者

來自一線程序員Seven的探索與實踐，持續(xù)學習迭代中~

本文已收錄于我的個人博客： https://www.seven97.top

公眾號：seven97，歡迎關注~