JDK集合源碼之解析TreeMap(一)

更新時間：2021年07月06日 10:42:00 作者：興趣使然的草帽路飛

下面小編就為大家?guī)硪黄獪\談java中的TreeMap 排序與TreeSet 排序。小編覺得挺不錯的，現(xiàn)在就分享給大家，也給大家做個參考。一起跟隨小編過來看看吧

簡介

TreeMap使用紅黑樹存儲元素，可以保證元素按key值的大小進(jìn)行遍歷。

繼承體系

TreeMap

TreeMap實現(xiàn)了Map、SortedMap、NavigableMap、Cloneable、Serializable等接口。

SortedMap規(guī)定了元素可以按key的大小來遍歷，它定義了一些返回部分map的方法。

public interface SortedMap<K,V> extends Map<K,V> {
    // key的比較器
    Comparator<? super K> comparator();
    // 返回fromKey（包含）到toKey（不包含）之間的元素組成的子map
    SortedMap<K,V> subMap(K fromKey, K toKey);
    // 返回小于toKey（不包含）的子map
    SortedMap<K,V> headMap(K toKey);
    // 返回大于等于fromKey（包含）的子map
    SortedMap<K,V> tailMap(K fromKey);
    // 返回最小的key
    K firstKey();
    // 返回最大的key
    K lastKey();
    // 返回key集合
    Set<K> keySet();
    // 返回value集合
    Collection<V> values();
    // 返回節(jié)點集合
    Set<Map.Entry<K, V>> entrySet();
}

NavigableMap是對SortedMap的增強(qiáng)，定義了一些返回離目標(biāo)key最近的元素的方法。

public interface NavigableMap<K,V> extends SortedMap<K,V> {
    // 小于給定key的最大節(jié)點
    Map.Entry<K,V> lowerEntry(K key);
    // 小于給定key的最大key
    K lowerKey(K key);
    // 小于等于給定key的最大節(jié)點
    Map.Entry<K,V> floorEntry(K key);
    // 小于等于給定key的最大key
    K floorKey(K key);
    // 大于等于給定key的最小節(jié)點
    Map.Entry<K,V> ceilingEntry(K key);
    // 大于等于給定key的最小key
    K ceilingKey(K key);
    // 大于給定key的最小節(jié)點
    Map.Entry<K,V> higherEntry(K key);
    // 大于給定key的最小key
    K higherKey(K key);
    // 最小的節(jié)點
    Map.Entry<K,V> firstEntry();
    // 最大的節(jié)點
    Map.Entry<K,V> lastEntry();
    // 彈出最小的節(jié)點
    Map.Entry<K,V> pollFirstEntry();
    // 彈出最大的節(jié)點
    Map.Entry<K,V> pollLastEntry();
    // 返回倒序的map
    NavigableMap<K,V> descendingMap();
    // 返回有序的key集合
    NavigableSet<K> navigableKeySet();
    // 返回倒序的key集合
    NavigableSet<K> descendingKeySet();
    // 返回從fromKey到toKey的子map，是否包含起止元素可以自己決定
    NavigableMap<K,V> subMap(K fromKey, boolean fromInclusive,
                             K toKey,   boolean toInclusive);
    // 返回小于toKey的子map，是否包含toKey自己決定
    NavigableMap<K,V> headMap(K toKey, boolean inclusive);
    // 返回大于fromKey的子map，是否包含fromKey自己決定
    NavigableMap<K,V> tailMap(K fromKey, boolean inclusive);
    // 等價于subMap(fromKey, true, toKey, false)
    SortedMap<K,V> subMap(K fromKey, K toKey);
    // 等價于headMap(toKey, false)
    SortedMap<K,V> headMap(K toKey);
    // 等價于tailMap(fromKey, true)
    SortedMap<K,V> tailMap(K fromKey);
}

存儲結(jié)構(gòu)

TreeMap-structure

TreeMap只使用到了紅黑樹，所以它的時間復(fù)雜度為O(log n)，我們再來回顧一下紅黑樹的特性。

（1）每個節(jié)點或者是黑色，或者是紅色。

（2）根節(jié)點是黑色。

（3）每個葉子節(jié)點（NIL）是黑色。（注意：這里葉子節(jié)點，是指為空(NIL或NULL)的葉子節(jié)點！）

（4）如果一個節(jié)點是紅色的，則它的子節(jié)點必須是黑色的。

（5）從一個節(jié)點到該節(jié)點的子孫節(jié)點的所有路徑上包含相同數(shù)目的黑節(jié)點。

源碼解析

屬性

/**
 * 比較器，如果沒傳則key要實現(xiàn)Comparable接口
 */
private final Comparator<? super K> comparator;
/**
 * 根節(jié)點
 */
private transient Entry<K,V> root;
/**
 * 元素個數(shù)
 */
private transient int size = 0;
/**
 * 修改次數(shù)
 */
private transient int modCount = 0;

（1）comparator

按key的大小排序有兩種方式，一種是key實現(xiàn)Comparable接口，一種方式通過構(gòu)造方法傳入比較器。

（2）root

根節(jié)點，TreeMap沒有桶的概念，所有的元素都存儲在一顆樹中。

Entry內(nèi)部類

存儲節(jié)點，典型的紅黑樹結(jié)構(gòu)。

static final class Entry<K,V> implements Map.Entry<K,V> {
    K key;
    V value;
    Entry<K,V> left;
    Entry<K,V> right;
    Entry<K,V> parent;
    boolean color = BLACK;
}

構(gòu)造方法

/**
 * 默認(rèn)構(gòu)造方法，key必須實現(xiàn)Comparable接口 
 */
public TreeMap() {
    comparator = null;
}
/**
 * 使用傳入的comparator比較兩個key的大小
 */
public TreeMap(Comparator<? super K> comparator) {
    this.comparator = comparator;
}
/**
 * key必須實現(xiàn)Comparable接口，把傳入map中的所有元素保存到新的TreeMap中 
 */
public TreeMap(Map<? extends K, ? extends V> m) {
    comparator = null;
    putAll(m);
}
/**
 * 使用傳入map的比較器，并把傳入map中的所有元素保存到新的TreeMap中 
 */
public TreeMap(SortedMap<K, ? extends V> m) {
    comparator = m.comparator();
    try {
        buildFromSorted(m.size(), m.entrySet().iterator(), null, null);
    } catch (java.io.IOException cannotHappen) {
    } catch (ClassNotFoundException cannotHappen) {
    }
}

構(gòu)造方法主要分成兩類，一類是使用comparator比較器，一類是key必須實現(xiàn)Comparable接口。

其實，筆者認(rèn)為這兩種比較方式可以合并成一種，當(dāng)沒有傳comparator的時候，可以用以下方式來給comparator賦值，這樣后續(xù)所有的比較操作都可以使用一樣的邏輯處理了，而不用每次都檢查comparator為空的時候又用Comparable來實現(xiàn)一遍邏輯。

// 如果comparator為空，則key必須實現(xiàn)Comparable接口，所以這里肯定可以強(qiáng)轉(zhuǎn)
// 這樣在構(gòu)造方法中統(tǒng)一替換掉，后續(xù)的邏輯就都一致了
comparator = (k1, k2) -> ((Comparable<? super K>)k1).compareTo(k2);

get(Object key)方法

獲取元素，典型的二叉查找樹的查找方法。

public V get(Object key) {
    // 根據(jù)key查找元素
    Entry<K,V> p = getEntry(key);
    // 找到了返回value值，沒找到返回null
    return (p==null ? null : p.value);
}
final Entry<K,V> getEntry(Object key) {
    // 如果comparator不為空，使用comparator的版本獲取元素
    if (comparator != null)
        return getEntryUsingComparator(key);
    // 如果key為空返回空指針異常
    if (key == null)
        throw new NullPointerException();
    // 將key強(qiáng)轉(zhuǎn)為Comparable
    @SuppressWarnings("unchecked")
    Comparable<? super K> k = (Comparable<? super K>) key;
    // 從根元素開始遍歷
    Entry<K,V> p = root;
    while (p != null) {
        int cmp = k.compareTo(p.key);
        if (cmp < 0)
            // 如果小于0從左子樹查找
            p = p.left;
        else if (cmp > 0)
            // 如果大于0從右子樹查找
            p = p.right;
        else
            // 如果相等說明找到了直接返回
            return p;
    }
    // 沒找到返回null
    return null;
}
final Entry<K,V> getEntryUsingComparator(Object key) {
    @SuppressWarnings("unchecked")
    K k = (K) key;
    Comparator<? super K> cpr = comparator;
    if (cpr != null) {
        // 從根元素開始遍歷
        Entry<K,V> p = root;
        while (p != null) {
            int cmp = cpr.compare(k, p.key);
            if (cmp < 0)
                // 如果小于0從左子樹查找
                p = p.left;
            else if (cmp > 0)
                // 如果大于0從右子樹查找
                p = p.right;
            else
                // 如果相等說明找到了直接返回
                return p;
        }
    }
    // 沒找到返回null
    return null;
}

（1）從root遍歷整個樹；

（2）如果待查找的key比當(dāng)前遍歷的key小，則在其左子樹中查找；

（3）如果待查找的key比當(dāng)前遍歷的key大，則在其右子樹中查找；

（4）如果待查找的key與當(dāng)前遍歷的key相等，則找到了該元素，直接返回；

（5）從這里可以看出是否有comparator分化成了兩個方法，但是內(nèi)部邏輯一模一樣，因此可見筆者comparator = (k1, k2) -> ((Comparable<? super K>)k1).compareTo(k2);這種改造的必要性。

特性再回顧

（1）每個節(jié)點或者是黑色，或者是紅色。

（2）根節(jié)點是黑色。

（3）每個葉子節(jié)點（NIL）是黑色。（注意：這里葉子節(jié)點，是指為空(NIL或NULL)的葉子節(jié)點?。?/p>

（4）如果一個節(jié)點是紅色的，則它的子節(jié)點必須是黑色的。

（5）從一個節(jié)點到該節(jié)點的子孫節(jié)點的所有路徑上包含相同數(shù)目的黑節(jié)點。

左旋

左旋，就是以某個節(jié)點為支點向左旋轉(zhuǎn)。

left-rotation

整個左旋過程如下：

（1）將 y的左節(jié)點設(shè)為 x的右節(jié)點，即將 β 設(shè)為 x的右節(jié)點；

（2）將 x 設(shè)為 y的左節(jié)點的父節(jié)點，即將 β的父節(jié)點設(shè)為 x；

（3）將 x的父節(jié)點設(shè)為 y的父節(jié)點；

（4）如果 x的父節(jié)點為空節(jié)點，則將y設(shè)置為根節(jié)點；如果x是它父節(jié)點的左（右）節(jié)點，則將y設(shè)置為x父節(jié)點的左（右）節(jié)點；

（5）將 x 設(shè)為 y的左節(jié)點；

（6）將 x的父節(jié)點設(shè)為 y；

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

/**
 * 以p為支點進(jìn)行左旋
 * 假設(shè)p為圖中的x
 */
private void rotateLeft(Entry<K,V> p) {
    if (p != null) {
        // p的右節(jié)點，即y
        Entry<K,V> r = p.right;
        // （1）將 y的左節(jié)點 設(shè)為 x的右節(jié)點
        p.right = r.left;
        // （2）將 x 設(shè)為 y的左節(jié)點的父節(jié)點（如果y的左節(jié)點存在的話）
        if (r.left != null)
            r.left.parent = p;
        // （3）將 x的父節(jié)點 設(shè)為 y的父節(jié)點
        r.parent = p.parent;
        // （4）...
        if (p.parent == null)
            // 如果 x的父節(jié)點 為空，則將y設(shè)置為根節(jié)點
            root = r;
        else if (p.parent.left == p)
            // 如果x是它父節(jié)點的左節(jié)點，則將y設(shè)置為x父節(jié)點的左節(jié)點
            p.parent.left = r;
        else
            // 如果x是它父節(jié)點的右節(jié)點，則將y設(shè)置為x父節(jié)點的右節(jié)點
            p.parent.right = r;
        // （5）將 x 設(shè)為 y的左節(jié)點
        r.left = p;
        // （6）將 x的父節(jié)點 設(shè)為 y
        p.parent = r;
    }
}

右旋

右旋，就是以某個節(jié)點為支點向右旋轉(zhuǎn)。

right-rotation

整個右旋過程如下：

（1）將 x的右節(jié)點設(shè)為 y的左節(jié)點，即將 β 設(shè)為 y的左節(jié)點；

（2）將 y 設(shè)為 x的右節(jié)點的父節(jié)點，即將 β的父節(jié)點設(shè)為 y；

（3）將 y的父節(jié)點設(shè)為 x的父節(jié)點；

（4）如果 y的父節(jié)點是空節(jié)點，則將x設(shè)為根節(jié)點；如果y是它父節(jié)點的左（右）節(jié)點，則將x設(shè)為y的父節(jié)點的左（右）節(jié)點；

（5）將 y 設(shè)為 x的右節(jié)點；

（6）將 y的父節(jié)點設(shè)為 x；

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

/**
 * 以p為支點進(jìn)行右旋
 * 假設(shè)p為圖中的y
 */
private void rotateRight(Entry<K,V> p) {
    if (p != null) {
        // p的左節(jié)點，即x
        Entry<K,V> l = p.left;
        // （1）將 x的右節(jié)點 設(shè)為 y的左節(jié)點
        p.left = l.right;
        // （2）將 y 設(shè)為 x的右節(jié)點的父節(jié)點（如果x有右節(jié)點的話）
        if (l.right != null) l.right.parent = p;
        // （3）將 y的父節(jié)點 設(shè)為 x的父節(jié)點
        l.parent = p.parent;
        // （4）...
        if (p.parent == null)
            // 如果 y的父節(jié)點 是 空節(jié)點，則將x設(shè)為根節(jié)點
            root = l;
        else if (p.parent.right == p)
            // 如果y是它父節(jié)點的右節(jié)點，則將x設(shè)為y的父節(jié)點的右節(jié)點
            p.parent.right = l;
        else
            // 如果y是它父節(jié)點的左節(jié)點，則將x設(shè)為y的父節(jié)點的左節(jié)點
            p.parent.left = l;
        // （5）將 y 設(shè)為 x的右節(jié)點
        l.right = p;
        // （6）將 y的父節(jié)點 設(shè)為 x
        p.parent = l;
    }
}

插入元素

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

public V put(K key, V value) {
    Entry<K,V> 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<K,V> parent;
    // 根據(jù)是否有comparator使用不同的分支
    Comparator<? super K> 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<? super K> k = (Comparable<? super K>) 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<K,V> e = new Entry<>(key, value, parent);
    if (cmp < 0)
        // 如果小于0插入到左子節(jié)點
        parent.left = e;
    else
        // 如果大于0插入到右子節(jié)點
        parent.right = e;
    // 插入之后的平衡
    fixAfterInsertion(e);
    // 元素個數(shù)加1（不需要擴(kuò)容）
    size++;
    // 修改次數(shù)加1
    modCount++;
    // 如果插入了新節(jié)點返回空
    return null;
}

插入再平衡

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

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

插入的元素如果是根節(jié)點，則直接涂成黑色即可，不用平衡；
插入的元素的父節(jié)點如果為黑色，不需要平衡；
插入的元素的父節(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<K,V> 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<K,V> 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<K,V> 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;
}