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

歡迎關(guān)注我的公眾號“彤哥讀源碼”，查看更多源碼系列文章, 與彤哥一起暢游源碼的海洋。

簡介

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

繼承體系

TreeMap實(shí)現(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é)點(diǎn)集合
    Set<Map.Entry<K, V>> entrySet();
}

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

public interface NavigableMap<K,V> extends SortedMap<K,V> {
    // 小于給定key的最大節(jié)點(diǎn)
    Map.Entry<K,V> lowerEntry(K key);
    // 小于給定key的最大key
    K lowerKey(K key);
    // 小于等于給定key的最大節(jié)點(diǎn)
    Map.Entry<K,V> floorEntry(K key);
    // 小于等于給定key的最大key
    K floorKey(K key);
    // 大于等于給定key的最小節(jié)點(diǎn)
    Map.Entry<K,V> ceilingEntry(K key);
    // 大于等于給定key的最小key
    K ceilingKey(K key);
    // 大于給定key的最小節(jié)點(diǎn)
    Map.Entry<K,V> higherEntry(K key);
    // 大于給定key的最小key
    K higherKey(K key);
    // 最小的節(jié)點(diǎn)
    Map.Entry<K,V> firstEntry();
    // 最大的節(jié)點(diǎn)
    Map.Entry<K,V> lastEntry();
    // 彈出最小的節(jié)點(diǎn)
    Map.Entry<K,V> pollFirstEntry();
    // 彈出最大的節(jié)點(diǎn)
    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只使用到了紅黑樹，所以它的時間復(fù)雜度為O(log n)，我們再來回顧一下紅黑樹的特性。

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

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

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

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

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

源碼解析

屬性

/**
 * 比較器，如果沒傳則key要實(shí)現(xiàn)Comparable接口
 */
private final Comparator<? super K> comparator;

/**
 * 根節(jié)點(diǎn)
 */
private transient Entry<K,V> root;

/**
 * 元素個數(shù)
 */
private transient int size = 0;

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

（1）comparator

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

（2）root

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

Entry內(nèi)部類

存儲節(jié)點(diǎn)，典型的紅黑樹結(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必須實(shí)現(xiàn)Comparable接口 
 */
public TreeMap() {
    comparator = null;
}

/**
 * 使用傳入的comparator比較兩個key的大小
 */
public TreeMap(Comparator<? super K> comparator) {
    this.comparator = comparator;
}

/**
 * key必須實(shí)現(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必須實(shí)現(xiàn)Comparable接口。

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

// 如果comparator為空，則key必須實(shí)現(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);這種改造的必要性。

我是一條美麗的分割線，前方高能，請做好準(zhǔn)備。

特性再回顧

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

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

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

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

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

左旋

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

整個左旋過程如下：

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

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

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

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

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

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

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

/**
 * 以p為支點(diǎn)進(jìn)行左旋
 * 假設(shè)p為圖中的x
 */
private void rotateLeft(Entry<K,V> p) {
    if (p != null) {
        // p的右節(jié)點(diǎn)，即y
        Entry<K,V> r = p.right;

        // （1）將 y的左節(jié)點(diǎn) 設(shè)為 x的右節(jié)點(diǎn)
        p.right = r.left;

        // （2）將 x 設(shè)為 y的左節(jié)點(diǎn)的父節(jié)點(diǎn)（如果y的左節(jié)點(diǎn)存在的話）
        if (r.left != null)
            r.left.parent = p;

        // （3）將 x的父節(jié)點(diǎn) 設(shè)為 y的父節(jié)點(diǎn)
        r.parent = p.parent;

        // （4）...
        if (p.parent == null)
            // 如果 x的父節(jié)點(diǎn) 為空，則將y設(shè)置為根節(jié)點(diǎn)
            root = r;
        else if (p.parent.left == p)
            // 如果x是它父節(jié)點(diǎn)的左節(jié)點(diǎn)，則將y設(shè)置為x父節(jié)點(diǎn)的左節(jié)點(diǎn)
            p.parent.left = r;
        else
            // 如果x是它父節(jié)點(diǎn)的右節(jié)點(diǎn)，則將y設(shè)置為x父節(jié)點(diǎn)的右節(jié)點(diǎn)
            p.parent.right = r;

        // （5）將 x 設(shè)為 y的左節(jié)點(diǎn)
        r.left = p;

        // （6）將 x的父節(jié)點(diǎn) 設(shè)為 y
        p.parent = r;
    }
}

右旋

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

整個右旋過程如下：

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

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

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

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

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

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

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

/**
 * 以p為支點(diǎn)進(jìn)行右旋
 * 假設(shè)p為圖中的y
 */
private void rotateRight(Entry<K,V> p) {
    if (p != null) {
        // p的左節(jié)點(diǎn)，即x
        Entry<K,V> l = p.left;

        // （1）將 x的右節(jié)點(diǎn) 設(shè)為 y的左節(jié)點(diǎn)
        p.left = l.right;

        // （2）將 y 設(shè)為 x的右節(jié)點(diǎn)的父節(jié)點(diǎn)（如果x有右節(jié)點(diǎn)的話）
        if (l.right != null) l.right.parent = p;

        // （3）將 y的父節(jié)點(diǎn) 設(shè)為 x的父節(jié)點(diǎn)
        l.parent = p.parent;

        // （4）...
        if (p.parent == null)
            // 如果 y的父節(jié)點(diǎn) 是 空節(jié)點(diǎn)，則將x設(shè)為根節(jié)點(diǎn)
            root = l;
        else if (p.parent.right == p)
            // 如果y是它父節(jié)點(diǎn)的右節(jié)點(diǎn)，則將x設(shè)為y的父節(jié)點(diǎn)的右節(jié)點(diǎn)
            p.parent.right = l;
        else
            // 如果y是它父節(jié)點(diǎn)的左節(jié)點(diǎn)，則將x設(shè)為y的父節(jié)點(diǎn)的左節(jié)點(diǎn)
            p.parent.left = l;

        // （5）將 y 設(shè)為 x的右節(jié)點(diǎn)
        l.right = p;

        // （6）將 y的父節(jié)點(diǎn) 設(shè)為 x
        p.parent = l;
    }
}

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

現(xiàn)在公眾號文章沒辦法留言了，如果有什么疑問或者建議請直接在公眾號給我留言。

歡迎關(guān)注我的公眾號“彤哥讀源碼”，查看更多源碼系列文章, 與彤哥一起暢游源碼的海洋。