乐于分享
好东西不私藏

JDK源码阅读(7) — TreeMap:红黑树在手,有序天下我有

JDK源码阅读(7) — TreeMap:红黑树在手,有序天下我有

JDK源码阅读(7) — TreeMap:红黑树在手,有序天下我有

一、类的基本介绍

HashMap 用哈希表实现了 O(1) 的增删改查,但它有一个硬伤——无序。当你需要按 key 有序遍历、范围查询、或者找最大最小 key 时,HashMap 就无能为力了。

这个时候 TreeMap 登场了:

public class TreeMap<K,V>
    extends AbstractMap<K,V>
    implements NavigableMap<K,V>, Cloneable, java.io.Serializable
  • 所在包:java.util
  • 底层结构:红黑树(Red-Black Tree)——一种自平衡的二叉查找树
  • 核心能力:所有操作 O(log n) 时间复杂度,key 有序
  • 排序依据:自然顺序(Comparable)或自定义 Comparator

跟 TreeSet 的关系?TreeSet 底层就是 TreeMap,value 用一个 PRESENT 占位对象,跟 HashSet 包装 HashMap 如出一辙。

二、核心源码解剖

2.1 节点结构

TreeMap 的每个节点就是一个 Entry,存储了 key-value 以及红黑树需要的指针和颜色标记:

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;  // 节点颜色:RED 或 BLACK

    Entry(K key, V value, Entry<K,V> parent) {
        this.key = key;
        this.value = value;
        this.parent = parent;
    }
}

2.2 全局变量

private transient Comparator<? super K> comparator;  // 比较器
private transient Entry<K,V> root;      // 红黑树根节点
private transient int size = 0;         // 节点数量
private transient int modCount = 0;     // 结构修改计数(Fast-Fail)

注意这里的 comparator —— 这是 TreeMap 排序的灵魂。如果你不传 Comparator,key 必须实现 Comparable 接口,否则 put 时会抛 ClassCastException。

2.3 构造方法

// 无参构造:使用自然顺序
public TreeMap() {
    comparator = null;
}

// 指定比较器
public TreeMap(Comparator<? super K> comparator) {
    this.comparator = comparator;
}

// 从已有 Map 构建
public TreeMap(Map<? extends K, ? extends V> m) {
    comparator = null;
    putAll(m);
}

// 从 SortedMap 构建(复用比较器)
public TreeMap(SortedMap<K, ? extends V> m) {
    comparator = m.comparator();
    try {
        buildFromSorted(m.size(), m.entrySet().iterator(), null, null);
    } catch (java.io.IOException | ClassNotFoundException cannotHappen) {
    }
}

最后这个 buildFromSorted 值得一提——它直接递归构建一棵平衡树,而不是逐个 put,这样初始化的时间复杂度从 O(n log n) 降到了 O(n)。JDK 源码处处是细节。

2.4 put() 方法——红黑树的灵魂

public V put(K key, V value) {
    Entry<K,V> t = root;
    // 1. 根节点为空,直接创建
    if (t == null) {
        compare(key, key); // type check
        root = new Entry<>(key, value, null);
        size = 1;
        modCount++;
        return null;
    }
    
    // 2. 二叉查找树插入
    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 {
        // 使用自然顺序(key 必须实现 Comparable)
        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);
    }
    
    // 3. 插入新节点(默认红色)
    Entry<K,V> e = new Entry<>(key, value, parent);
    if (cmp < 0)
        parent.left = e;
    else
        parent.right = e;
    
    // 4. 红黑树自平衡(关键!)
    fixAfterInsertion(e);
    size++;
    modCount++;
    return null;
}

插入逻辑很清晰:BST 查找 → 定位父节点 → 插入红色节点 → 红黑树修正。

2.5 fixAfterInsertion——红黑树修正

红黑树有 5 条铁律:

  1. 每个节点是红色或黑色
  2. 根节点是黑色
  3. 叶子节点(NIL)是黑色
  4. 红色节点不能有红色孩子(不能连续红色)
  5. 任一节点到每个叶子的路径上黑色节点数相同

插入红色节点后可能违反规则 4,需要修复:

private void fixAfterInsertion(Entry<K,V> x) {
    x.color = RED;
    
    while (x != null && x != root && x.parent.color == RED) {
        if (parentOf(x) == leftOf(parentOf(parentOf(x)))) {
            Entry<K,V> y = rightOf(parentOf(parentOf(x))); // 叔叔节点
            if (colorOf(y) == RED) {
                // Case 1: 叔叔是红色 → 颜色翻转
                setColor(parentOf(x), BLACK);
                setColor(y, BLACK);
                setColor(parentOf(parentOf(x)), RED);
                x = parentOf(parentOf(x));
            } else {
                if (x == rightOf(parentOf(x))) {
                    // Case 2: 叔叔是黑色,x 是右孩子 → 左旋
                    x = parentOf(x);
                    rotateLeft(x);
                }
                // Case 3: 叔叔是黑色,x 是左孩子 → 右旋
                setColor(parentOf(x), BLACK);
                setColor(parentOf(parentOf(x)), RED);
                rotateRight(parentOf(parentOf(x)));
            }
        } else {
            // 对称情况(父节点是右孩子)
            // ...
        }
    }
    root.color = BLACK;
}

三种情况的处理可以概括为一张流程图:

插入修复三态:

  • 👉 叔叔是红色 → 向上传递红色(把祖父变红,父和叔变黑)
  • 👉 叔叔是黑色,当前是内侧孩子 → 先旋转变为外侧
  • 👉 叔叔是黑色,当前是外侧孩子 → 旋转父节点+祖父节点变色

2.6 get() 方法——O(log n) 查找

public V get(Object key) {
    Entry<K,V> p = getEntry(key);
    return (p == null ? null : p.value);
}

final Entry<K,V> getEntry(Object key) {
    // 比较器路径
    if (comparator != null)
        return getEntryUsingComparator(key);
    if (key == null)
        throw new NullPointerException();
    
    // 自然顺序路径
    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;
}

经典 BST 查找,没有哈希表的魔法,但胜在稳定——每次查询都是确确实实的 O(log n)。

三、NavigableMap 的花式操作

TreeMap 实现了 NavigableMap 接口,提供了很多"有序"场景下的实用方法:

// 范围查询
K firstKey();                    // 最小 key
K lastKey();                     // 最大 key
K lowerKey(K key);               // 严格小于 key 的最大 key
K floorKey(K key);               // 小于等于 key 的最大 key
K ceilingKey(K key);             // 大于等于 key 的最小 key
K higherKey(K key);              // 严格大于 key 的最小 key

// 子视图
SortedMap<K,V> subMap(K fromKey, K toKey);         // [from, to)
SortedMap<K,V> headMap(K toKey);                   // [最小, to)
SortedMap<K,V> tailMap(K fromKey);                 // [from, 最大)
NavigableMap<K,V> descendingMap();                 // 倒序视图

这些方法的实现都非常巧妙,比如 lowerKey 就是在红黑树上做一次"寻找前驱"的遍历:

final Entry<K,V> getLowerEntry(K key) {
    Entry<K,V> p = root;
    while (p != null) {
        int cmp = compare(key, p.key);
        if (cmp > 0) {
            // key > p.key,p 可能是候选,继续往右找更大的
            if (p.right != null)
                p = p.right;
            else
                return p;
        } else {
            // key <= p.key,往左找
            if (p.left != null)
                p = p.left;
            else {
                // 左子树空了,回溯找祖先中第一个"从左侧来的"
                Entry<K,V> parent = p.parent;
                Entry<K,V> ch = p;
                while (parent != null && ch == parent.left) {
                    ch = parent;
                    parent = parent.parent;
                }
                return parent;
            }
        }
    }
    return null;
}

四、完整使用示例

import java.util.*;

public class TreeMapDemo {
    public static void main(String[] args) {
        // 1. 基本用法:按字符串自然顺序(字典序)
        TreeMap<String, Integer> scores = new TreeMap<>();
        scores.put("Alice", 95);
        scores.put("Bob", 87);
        scores.put("Charlie", 92);
        scores.put("David", 78);
        
        System.out.println("按姓名排序的成绩单:");
        for (Map.Entry<String, Integer> entry : scores.entrySet()) {
            System.out.println(entry.getKey() + ": " + entry.getValue());
        }
        // 输出: Alice:95 → Bob:87 → Charlie:92 → David:78(自然顺序)
        
        // 2. 范围查询
        System.out.println(" C开头的学生:");
        System.out.println(scores.subMap("C", "D"));
        // 输出: {Charlie=92}
        
        System.out.println("A到C之间的学生:");
        System.out.println(scores.subMap("A", "C"));  // [A, C)
        // 输出: {Alice=95, Bob=87}
        
        // 3. 边界查询
        System.out.println(" 最高分:" + scores.lastEntry());  // David=78
        System.out.println("最低分:" + scores.firstEntry());  // Alice=95
        System.out.println("比Bob大的最小key:" + scores.higherKey("Bob"));  // Charlie
        System.out.println("Bob或之前的学生:" + scores.headMap("Bob", true));
        
        // 4. 自定义排序:按分数倒序
        TreeMap<Integer, String> byScore = new TreeMap<>(Comparator.reverseOrder());
        byScore.put(95, "Alice");
        byScore.put(87, "Bob");
        byScore.put(92, "Charlie");
        byScore.put(78, "David");
        
        System.out.println(" 按分数从高到低:");
        for (Map.Entry<Integer, String> entry : byScore.entrySet()) {
            System.out.println(entry.getValue() + ": " + entry.getKey());
        }
        // 输出: Alice:95 → Charlie:92 → Bob:87 → David:78
        
        // 5. 倒序遍历
        System.out.println(" 倒序遍历原map:");
        NavigableMap<String, Integer> desc = scores.descendingMap();
        desc.forEach((k, v) -> System.out.println(k + ": " + v));
    }
}

五、TreeMap vs HashMap 深度对比

维度 HashMap TreeMap
底层结构 哈希表(数组+链表/红黑树) 红黑树
时间复杂度 O(1) 平均,O(log n) 最坏 O(log n) 保证
有序性 无序 有序(自然顺序或 Comparator)
null key 允许一个 null 不允许(compareTo 会抛 NPE)
null value 允许 允许
内存占用 较低(数组+节点) 较高(每个节点有 parent/left/right/color 字段)
范围查询 ❌ 不支持 ✅ subMap/headMap/tailMap
初始化 默认 16 个桶 空树,插入创建根节点
迭代顺序 不确定(可能变,rehash 后) 稳定(key 排序)
适用场景 通用键值存储,查找为主 需要有序遍历、范围查询

六、注意事项 & 大坑

⚠️ 1. TreeMap 不是线程安全的

多线程环境下需要加锁或使用 Collections.synchronizedSortedMap(new TreeMap<>())。并发场景可以考虑 ConcurrentSkipListMap。

⚠️ 2. 比较器必须与 equals 一致

跟 TreeSet 一样,TreeMap 用 Comparator.compare(a, b) == 0 而不是 equals() 判断 key 是否相等。如果 compare 返回 0 但 equals 返回 false,行为是未定义的——这会导致数据"丢失"。

⚠️ 3. key 不能为 null

TreeMap 的 put(null, value) 会抛 NullPointerException,因为 compare(null, key) 会调用 null.compareTo(...) 或 comparator.compare(null, key)。而 HashMap 允许一个 null key。

⚠️ 4. 迭代器是 fail-fast 的

跟所有集合类一样,迭代过程中如果其他线程修改了 TreeMap,会抛 ConcurrentModificationException。

⚠️ 5. subMap 返回的是视图,不是快照

TreeMap<String, Integer> map = new TreeMap<>();
map.put("a", 1);
map.put("b", 2);
map.put("c", 3);
SortedMap<String, Integer> sub = map.subMap("a", "c");
sub.put("b", 99);   // ✅ 可以修改值
map.put("ab", 15);  // ✅ 原 map 插入在范围内,子视图也会看到
map.put("d", 4);    // ❌ 子视图不变(d 不在范围内)

这个特性很强大但也容易踩坑——不要在迭代子视图时修改原 map 的 key 顺序。

七、红黑树 VS 跳跃表

Java 里有两个有序 Map 实现:TreeMap(红黑树)和 ConcurrentSkipListMap(跳跃表)。

特性 红黑树 (TreeMap) 跳跃表 (ConcurrentSkipListMap)
时间复杂度 O(log n) O(log n) 期望
并发友好 ❌ 需要全表锁 ✅ 无锁并发(CAS)
空间占用 3 指针 + 1 颜色 平均 2 指针(每个节点有 1/n 的概率多一层)
实现复杂度 复杂(旋转、颜色翻转) 相对简单
范围查询 需中序遍历 可通过 forward 指针高效遍历

简单总结:单线程用 TreeMap,并发用 ConcurrentSkipListMap。

八、总结

TreeMap 是 JDK 中红黑树的教科书级实现,它能让你:

  1. ✅ 按 key 有序遍历 —— 自然顺序或自定义顺序
  2. ✅ 范围查询 —— subMap / headMap / tailMap 三件套
  3. ✅ 边界操作 —— 找最大最小、上下界
  4. ✅ 倒序遍历 —— descendingMap / descendingKeySet
  5. ✅ 稳定的 O(log n) —— 没有哈希冲突的烦恼

但也有限制:不允许 null key、非线程安全、内存开销大于 HashMap。

下一篇预告: ReentrantReadWriteLock —— 读写分离的锁哲学,跟 ReentrantLock 对比着学。

相关学习资料