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

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

簡介

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

繼承體系

TreeMap實現了Map、SortedMap、NavigableMap、Cloneable、Serializable等接口。

SortedMap規定了元素可以按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();
    // 返回節點集合
    Set<Map.Entry<K, V>> entrySet();
}

NavigableMap是對SortedMap的增強，定義了一些返回離目標key最近的元素的方法。

public interface NavigableMap<K,V> extends SortedMap<K,V> {
    // 小于給定key的最大節點
    Map.Entry<K,V> lowerEntry(K key);
    // 小于給定key的最大key
    K lowerKey(K key);
    // 小于等于給定key的最大節點
    Map.Entry<K,V> floorEntry(K key);
    // 小于等于給定key的最大key
    K floorKey(K key);
    // 大于等于給定key的最小節點
    Map.Entry<K,V> ceilingEntry(K key);
    // 大于等于給定key的最小key
    K ceilingKey(K key);
    // 大于給定key的最小節點
    Map.Entry<K,V> higherEntry(K key);
    // 大于給定key的最小key
    K higherKey(K key);
    // 最小的節點
    Map.Entry<K,V> firstEntry();
    // 最大的節點
    Map.Entry<K,V> lastEntry();
    // 彈出最小的節點
    Map.Entry<K,V> pollFirstEntry();
    // 彈出最大的節點
    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);
}

存儲結構

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

（1）每個節點或者是黑色，或者是紅色。

（2）根節點是黑色。

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

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

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

源碼解析

屬性

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

/**
 * 根節點
 */
private transient Entry<K,V> root;

/**
 * 元素個數
 */
private transient int size = 0;

/**
 * 修改次數
 */
private transient int modCount = 0;

（1）comparator

按key的大小排序有兩種方式，一種是key實現Comparable接口，一種方式通過構造方法傳入比較器。

（2）root

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

Entry內部類

存儲節點，典型的紅黑樹結構。

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;
}

構造方法

/**
 * 默認構造方法，key必須實現Comparable接口 
 */
public TreeMap() {
    comparator = null;
}

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

/**
 * key必須實現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) {
    }
}

構造方法主要分成兩類，一類是使用comparator比較器，一類是key必須實現Comparable接口。

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

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

get(Object key)方法

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

public V get(Object key) {
    // 根據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強轉為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比當前遍歷的key小，則在其左子樹中查找；

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

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

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

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

特性再回顧

（1）每個節點或者是黑色，或者是紅色。

（2）根節點是黑色。

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

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

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

左旋

左旋，就是以某個節點為支點向左旋轉。

整個左旋過程如下：

（1）將 y的左節點 設為 x的右節點，即將 β 設為 x的右節點；

（2）將 x 設為 y的左節點的父節點，即將 β的父節點 設為 x；

（3）將 x的父節點 設為 y的父節點；

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

（5）將 x 設為 y的左節點；

（6）將 x的父節點 設為 y；

讓我們來看看TreeMap中的實現：

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

        // （1）將 y的左節點 設為 x的右節點
        p.right = r.left;

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

        // （3）將 x的父節點 設為 y的父節點
        r.parent = p.parent;

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

        // （5）將 x 設為 y的左節點
        r.left = p;

        // （6）將 x的父節點 設為 y
        p.parent = r;
    }
}

右旋

右旋，就是以某個節點為支點向右旋轉。

整個右旋過程如下：

（1）將 x的右節點 設為 y的左節點，即 將 β 設為 y的左節點；

（2）將 y 設為 x的右節點的父節點，即 將 β的父節點 設為 y；

（3）將 y的父節點 設為 x的父節點；

（4）如果 y的父節點 是 空節點，則將x設為根節點；如果y是它父節點的左（右）節點，則將x設為y的父節點的左（右）節點；

（5）將 y 設為 x的右節點；

（6）將 y的父節點 設為 x；

讓我們來看看TreeMap中的實現：

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

        // （1）將 x的右節點 設為 y的左節點
        p.left = l.right;

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

        // （3）將 y的父節點 設為 x的父節點
        l.parent = p.parent;

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

        // （5）將 y 設為 x的右節點
        l.right = p;

        // （6）將 y的父節點 設為 x
        p.parent = l;
    }
}

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

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

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