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 条铁律:
每个节点是红色或黑色 根节点是黑色 叶子节点(NIL)是黑色 红色节点不能有红色孩子(不能连续红色) 任一节点到每个叶子的路径上黑色节点数相同
插入红色节点后可能违反规则 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 中红黑树的教科书级实现,它能让你:
✅ 按 key 有序遍历 —— 自然顺序或自定义顺序 ✅ 范围查询 —— subMap / headMap / tailMap 三件套 ✅ 边界操作 —— 找最大最小、上下界 ✅ 倒序遍历 —— descendingMap / descendingKeySet ✅ 稳定的 O(log n) —— 没有哈希冲突的烦恼
但也有限制:不允许 null key、非线程安全、内存开销大于 HashMap。
下一篇预告: ReentrantReadWriteLock —— 读写分离的锁哲学,跟 ReentrantLock 对比着学。
夜雨聆风