TreeSet 與 TreeMap 淺析：紅黑樹底層原理、排序機制與源碼解讀

TreeSet 與 TreeMap 淺析：紅黑樹底層原理、排序機制與源碼解讀

TreeSet 與 TreeMap 是 Java 集合框架中有序容器的核心實現。本文揭示兩者的包裝關係，深入剖析 TreeMap.put() 插入邏輯、getEntry() 查找邏輯，並以圖解方式演示紅黑樹的中序遍歷如何產生有序序列。

一、TreeSet 與 TreeMap 的包裝關係

TreeSet 內部持有一個 NavigableMap<E, Object> 變數 map，實際類別為 TreeMap。所有操作委派給該 Map：

// TreeSet 關鍵片段
private transient NavigableMap<E, Object> map;
private static final Object PRESENT = new Object();

public TreeSet() {
    this.map = new TreeMap<>();
}

public boolean add(E e) {
    return map.put(e, PRESENT) == null;
}

因此，掌握 TreeMap 即掌握 TreeSet。

二、TreeMap.put() 源碼深度解析

核心邏輯：從根節點開始，根據比較器（自然排序或定制排序）決定往左/右子樹，找到插入點建立 Entry 節點，最後修復紅黑樹性質。

public V put(K key, V value) {
    Entry<K,V> t = root;
    if (t == null) {                    // 空樹：建立根節點
        root = new Entry<>(key, value, null);
        size = 1;
        modCount++;
        return null;
    }
    int cmp;
    Entry<K,V> parent;
    Comparator<? super K> cpr = comparator;
    if (cpr != null) {                  // 定制排序
        do {
            parent = t;
            cmp = cpr.compare(key, t.key);
            if (cmp < 0) t = t.left;
            else if (cmp > 0) t = t.right;
            else return t.setValue(value); // key 重複，覆蓋值
        } while (t != null);
    } else {                            // 自然排序
        if (key == null) throw new NullPointerException();
        @SuppressWarnings("unchecked")
        Comparable<? super K> k = (Comparable<? super K>) 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==null，parent 是插入點的父節點
    Entry<K,V> e = new Entry<>(key, value, parent);
    if (cmp < 0) parent.left = e;
    else parent.right = e;
    fixAfterInsertion(e);               // 紅黑樹修復：旋轉+變色
    size++;
    modCount++;
    return null;
}

關鍵點：

1. 比較器非空用 cpr.compare，為空用 Comparable.compareTo。
2. key 為 null 自然排序拋 NPE；定制排序允許 null（視比較器實現）。
3. 插入後必須呼叫 fixAfterInsertion 維持紅黑樹 5 大性質。

三、查找邏輯：getEntry() 與 getEntryUsingComparator()

// 自然排序查找
final Entry<K,V> getEntry(Object key) {
    if (comparator != null) return getEntryUsingComparator(key);
    if (key == null) throw new NullPointerException();
    @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) p = p.left;
        else if (cmp > 0) p = p.right;
        else return p;                  // 找到
    }
    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) p = p.left;
            else if (cmp > 0) p = p.right;
            else return p;
        }
    }
    return null;
}

查找同樣是二分搜尋，時間複雜度 O(log n)。

四、紅黑樹中序遍歷產生有序序列

TreeMap 實作 NavigableMap，其 keySet()、values()、entrySet() 迭代器皆按中序遍歷（左-根-右）輸出，保證升序。

遍歷邏輯（簡化）：

1. 從根節點出發，沿左子樹一直走到最左葉節點（最小鍵）。
2. 輸出該節點。
3. 若有右子樹，轉向右子樹的最左葉節點；否則回溯至最近的「左轉祖先」。
4. 重複直到遍歷完所有節點。

圖中數字標示遍歷順序：1→2→3→4→5→6→7，對應鍵值由小到大。

五、TreeSet/TreeMap 排序總結

六、常見效能陷阱

• 比較器不一致：compare 回傳 0 但 equals 回傳 false，導致 TreeSet 視為重複而 HashSet 視為不同。
• 大量插入：O(log n) 慢於 HashMap O(1)。若不需有序，優先用 HashMap/HashSet。
• 併發：非執行緒安全。高併發用 ConcurrentSkipListMap/Set（跳表實現，同樣有序）。

TreeMap 以紅黑樹為骨架，TreeSet 輕量包裝，兩者共享同一套插入、查找、遍歷邏輯。讀懂 put()、getEntry()、fixAfterInsertion() 三大核心，即可洞悉有序容器的本質。