TreeMap
TreeMap 源码分析
TreeMap 解决的问题:既要 key-value 映射,又要 key 全局有序、能做范围查询(floor/ceiling/subMap)。
它用一棵红黑树把增删查全部压在稳定 O(log n),代价是 key 必须可比较(实现 Comparable 或传入 Comparator)。
HashMap 里的红黑树只是长链退化的兜底、结构不完整,想吃透红黑树的旋转与变色,看 TreeMap 这份完整实现。
// 基于 JDK 25 (本地 D:/1ForCode/JAVA_Source), java.util.TreeMap —— 结构与查找
public class TreeMap<K,V>
extends AbstractMap<K,V>
implements NavigableMap<K,V>, Cloneable, java.io.Serializable
{
@SuppressWarnings("serial") // Conditionally serializable
private final Comparator<? super K> comparator; // null 就用 key 自身的 Comparable，二选一决定全树的序
private transient Entry<K,V> root; // 整棵红黑树只握一个根引用
private transient int size = 0;
private transient int modCount = 0; // 结构修改计数，迭代器 fail-fast 靠它
private static final boolean RED = false;
private static final boolean BLACK = true;
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; // 有 parent 指针，旋转和找前驱后继才能不用递归/栈
boolean color = BLACK; // 默认黑；插入时由 fixAfterInsertion 第一句改染红
Entry(K key, V value, Entry<K,V> parent) {
this.key = key;
this.value = value;
this.parent = parent;
}
// ... getKey/getValue/setValue/equals/hashCode/toString 略
}
// 全 TreeMap 的比较入口：comparator 优先，否则强转 Comparable。key 类型不可比时在这里抛 ClassCastException
@SuppressWarnings("unchecked")
final int compare(Object k1, Object k2) {
return comparator==null ? ((Comparable<? super K>)k1).compareTo((K)k2)
: comparator.compare((K)k1, (K)k2);
}
// get 的核心：标准 BST 下降，红黑树保证树高 O(log n)
final Entry<K,V> getEntry(Object key) {
// Offload comparator-based version for sake of performance
if (comparator != null)
return getEntryUsingComparator(key); // comparator 版拆成独立方法，热路径循环里少一次分支
Objects.requireNonNull(key); // 自然排序下 null 没法 compareTo，所以 TreeMap 禁 null key
@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; // 坑：命中只看 compare == 0，与 equals 无关
}
return null;
}
}// 基于 JDK 25 (本地 D:/1ForCode/JAVA_Source), java.util.TreeMap —— 插入与红黑修复
public V put(K key, V value) {
return put(key, value, true); // put 和 putIfAbsent 共用一次树下降，只差 replaceOld 开关
}
@Override
public V putIfAbsent(K key, V value) {
return put(key, value, false);
}
private V put(K key, V value, boolean replaceOld) {
Entry<K,V> t = root;
if (t == null) {
addEntryToEmptyMap(key, value); // 空树单独走：里面 compare(key, key) 提前做类型/null 检查
return null;
}
int cmp;
Entry<K,V> parent;
// split comparator and comparable paths
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 {
V oldValue = t.value;
if (replaceOld || oldValue == null) {
t.value = value; // key 已存在：只覆盖 value，不动树、不加 modCount
}
return oldValue;
}
} while (t != null);
} else {
// ... Comparable 分支：除了换成 k.compareTo(t.key)，逐行相同
}
addEntry(key, value, parent, cmp < 0); // 下降到 null 才是真插入；cmp 是最后一次比较，决定挂左还是右
return null;
}
private void addEntry(K key, V value, Entry<K, V> parent, boolean addToLeft) {
Entry<K,V> e = new Entry<>(key, value, parent);
if (addToLeft)
parent.left = e;
else
parent.right = e;
fixAfterInsertion(e); // 先按 BST 挂好，再修红黑性质：定位与平衡两步解耦
size++;
modCount++;
}
private void addEntryToEmptyMap(K key, V value) {
compare(key, key); // type (and possibly null) check
root = new Entry<>(key, value, null);
size = 1;
modCount++;
}
private static <K,V> boolean colorOf(Entry<K,V> p) {
return (p == null ? BLACK : p.color); // null 视为黑：用 null-safe 访问器替代 CLR 教科书里的 nil 哨兵节点
}
// ... parentOf/leftOf/rightOf/setColor 同样做了 null 保护，略
/** From CLR */
private void rotateLeft(Entry<K,V> p) {
if (p != null) {
Entry<K,V> r = p.right;
p.right = r.left; // r 的左子树过继给 p：整个旋转唯一换爹的子树
if (r.left != null)
r.left.parent = p;
r.parent = p.parent; // r 顶替 p 接到祖父下面
if (p.parent == null)
root = r;
else if (p.parent.left == p)
p.parent.left = r;
else
p.parent.right = r;
r.left = p; // p 降级为 r 的左孩子。全程不改中序顺序，只调树形
p.parent = r;
}
}
// rotateRight 与之完全镜像（left/right 对调），略
/** From CLR */
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))); // y = 叔叔
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 在内侧：先旋成外侧（LR 转 LL）
x = parentOf(x);
rotateLeft(x);
}
setColor(parentOf(x), BLACK); // case 3 外侧：变色 + 对祖父一次旋转，循环就此终止
setColor(parentOf(parentOf(x)), RED);
rotateRight(parentOf(parentOf(x)));
}
} else {
// ... 父节点是祖父右孩子的镜像分支，逻辑对称，略
}
}
root.color = BLACK; // 根恒黑：case 1 把红一路推到根时在这里收口
}// 基于 JDK 25 (本地 D:/1ForCode/JAVA_Source), java.util.TreeMap —— 删除与范围查询
static <K,V> TreeMap.Entry<K,V> successor(Entry<K,V> t) {
if (t == null)
return null;
else if (t.right != null) { // 有右子树：后继 = 右子树最左节点
Entry<K,V> p = t.right;
while (p.left != null)
p = p.left;
return p;
} else { // 无右子树：往上找第一个"把自己当左子树"的祖先
Entry<K,V> p = t.parent;
Entry<K,V> ch = t;
while (p != null && ch == p.right) {
ch = p;
p = p.parent;
}
return p; // 迭代器 next() 也走这里：靠 parent 指针中序遍历，无需栈
}
}
private void deleteEntry(Entry<K,V> p) {
modCount++;
size--;
// If strictly internal, copy successor's element to p and then make p
// point to successor.
if (p.left != null && p.right != null) {
Entry<K,V> s = successor(p);
p.key = s.key; // 双孩子不挪节点：拷贝后继的 k/v，改删后继（后继至多一个右孩子）
p.value = s.value;
p = s;
} // p has 2 children
// Start fixup at replacement node, if it exists.
Entry<K,V> replacement = (p.left != null ? p.left : p.right);
if (replacement != null) {
// Link replacement to parent
replacement.parent = p.parent;
if (p.parent == null)
root = replacement;
else if (p == p.parent.left)
p.parent.left = replacement;
else
p.parent.right = replacement;
// Null out links so they are OK to use by fixAfterDeletion.
p.left = p.right = p.parent = null;
// Fix replacement
if (p.color == BLACK)
fixAfterDeletion(replacement); // 只有删黑节点才破坏黑高；删红节点白删，不用修
} else if (p.parent == null) { // return if we are the only node.
root = null;
} else { // No children. Use self as phantom replacement and unlink.
if (p.color == BLACK) // 叶子：先拿自己当"幽灵替身"原地修完，再从树上摘链
fixAfterDeletion(p);
if (p.parent != null) {
if (p == p.parent.left)
p.parent.left = null;
else if (p == p.parent.right)
p.parent.right = null;
p.parent = null;
}
}
}
/** From CLR */
private void fixAfterDeletion(Entry<K,V> x) {
while (x != root && colorOf(x) == BLACK) { // x 背着"双重黑"，循环目标是把多出的一重黑消化掉
if (x == leftOf(parentOf(x))) {
Entry<K,V> sib = rightOf(parentOf(x));
if (colorOf(sib) == RED) { // 兄弟红：旋转换出一个黑兄弟，归约成下面的情况
setColor(sib, BLACK);
setColor(parentOf(x), RED);
rotateLeft(parentOf(x));
sib = rightOf(parentOf(x));
}
if (colorOf(leftOf(sib)) == BLACK &&
colorOf(rightOf(sib)) == BLACK) { // 兄弟俩孩子全黑：兄弟染红，双重黑上移给父节点
setColor(sib, RED);
x = parentOf(x);
} else {
if (colorOf(rightOf(sib)) == BLACK) { // 近侄红远侄黑：先旋成"远侄红"
setColor(leftOf(sib), BLACK);
setColor(sib, RED);
rotateRight(sib);
sib = rightOf(parentOf(x));
}
setColor(sib, colorOf(parentOf(x))); // 远侄红：一次旋转+变色彻底解决
setColor(parentOf(x), BLACK);
setColor(rightOf(sib), BLACK);
rotateLeft(parentOf(x));
x = root; // 置 root 直接跳出循环
}
} else { // symmetric
// ... x 是右孩子的镜像分支，略
}
}
setColor(x, BLACK); // 出口补一刀黑：把"红+黑"落地成纯黑，或消掉根上的双重黑
}
// NavigableMap 的招牌能力：范围查询。BST 天然支持，HashMap 给不了
final Entry<K,V> getFloorEntry(K key) {
Entry<K,V> p = root;
while (p != null) {
int cmp = compare(key, p.key);
if (cmp > 0) {
if (p.right != null)
p = p.right;
else
return p; // 右边没有更大的了：当前节点就是 <= key 的最大者
} else if (cmp < 0) {
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;
}
} else
return p;
}
return null;
}
public Map.Entry<K,V> floorEntry(K key) {
return exportEntry(getFloorEntry(key)); // 导出 SimpleImmutableEntry 快照，防调用方持有内部节点乱改
}原理串讲
以 map.put(42, v) 走一遍完整链路。入口 put(K,V) 直接转发 put(key, value, true),这个私有三参版本是 put 和 putIfAbsent 的公共实现:树下降的代码只写一份,replaceOld 开关决定命中已有 key 时覆不覆盖。
进来先判空树,空树走 addEntryToEmptyMap,其中那句看似自恋的 compare(key, key) 是故意的——树里还没有别的 key 可比,就拿 key 和自己比一次,把 ClassCastException 和 NullPointerException 提前到插入第一个元素时抛出,而不是留到第二次 put 才炸。
非空树则进入 do-while 下降:每一步用 cpr.compare(或 compareTo)决定往左还是往右,同时用 parent 变量记住脚下的节点。
比出 0 说明 key 已存在,改个 value 就返回,连 modCount 都不动,因为没有发生结构修改,正在遍历的迭代器不该因此 fail-fast。
下降到 null 才是真插入,addEntry 按最后一次 cmp 的符号把新 Entry 挂到 parent 左或右,然后进入本类的灵魂 fixAfterInsertion。
为什么新节点要先染红?因为红黑树五条性质里最贵的一条是”任一节点到叶子的黑节点数相等”(黑高一致)。
代码块收起展开
挂一个黑节点必然让这条路径黑高 +1,全树性质大面积崩塌;挂红节点则黑高完全不变,最坏只违反"红节点的孩子必须黑"这一条局部性质,而且冲突点就在新节点和它父亲之间——破坏面最小、修复起点明确。
修复循环只处理"父亲也是红"的情况,按叔叔颜色二分:叔叔红就纯变色(父叔变黑、祖父变红),把冲突整体上移两层继续循环,一次旋转都不做;叔叔黑就先把"内侧孩子"旋成"外侧"(case 2),再对祖父做一次旋转加变色(case 3),此时性质完全恢复,循环退出。所以插入最多两次旋转,变色可能一路传到根,最后 root.color = BLACK 兜底。
旋转本身(rotateLeft/rotateRight)只重排局部指针,中序遍历顺序不变,也就是”有序”这个大前提在任何时刻都不被破坏。
删除走 deleteEntry。最巧的一步是双孩子的情况:不真的把那个节点从树里摘下来,而是找 successor(右子树最左节点),把后继的 key/value 拷进当前节点,然后转头去删后继。
为什么这么设计?因为后继必然没有左孩子,删除问题从”双孩子”降维成”至多一个孩子”,只剩三种平凡情形;而拷贝 k/v 不改变任何指针和颜色,树结构毫发无损。
代价是 Entry 对象和 key 的”身份”发生了偷换,这也是源码注释特意标注的原因。真正的黑高修复在 fixAfterDeletion:只有被摘的节点是黑色才需要修(删红不影响黑高),修复模型是让替身节点背上”双重黑”,看兄弟节点的颜色分四种情况,要么把双重黑上推,要么靠旋转从兄弟那边借一个黑节点过来,最多三次旋转终结。
还有一个工程细节:CLR 教科书算法依赖一个全局 nil 哨兵节点来统一处理空指针,JDK 不想为每棵树多存哨兵、也不想让代码到处判空,就用 colorOf/parentOf/leftOf/rightOf/setColor 这组 null-safe 静态访问器代替——colorOf(null) 返回 BLACK 恰好和”叶子(null)视为黑”的定义吻合。
查询侧,getEntry 就是纯 BST 下降;getFloorEntry 这类导航方法多一层回溯逻辑:向左走到头还没找到 <= key 的节点时,沿 parent 指针爬回第一个”把自己当右子树”的祖先。
注意 TreeMap 的”相等”完全由 compare == 0 定义,与 equals 无关——这既是它不允许 null key 的原因(null 没法参与比较),也是 comparator 与 equals 不一致时 Map 契约被打破的根源。
设计取舍
- 红黑树是”弱平衡”:只保证最长路径不超过最短路径 2 倍,换来插入最多 2 次、删除最多 3 次旋转,比 AVL 的严格平衡更适合写操作多的场景,所以 TreeMap、HashMap 树化桶、Linux CFS 调度器都选它。
- 判等只认
compare == 0不认 equals:自定义 Comparator 只比对象的某个字段时,两个 equals 不同的对象会被当成同一个 key 互相覆盖,这是线上常踩的坑。 - 不允许 null key(value 可以 null):key 要参与比较,
null.compareTo直接 NPE,源码在getEntry和addEntryToEmptyMap里都做了前置检查。 - 选型:只要 key、随机存取用 HashMap(平均 O(1) 但无序);要按 key 有序遍历或范围查询(如”找 100~200 之间的所有 key”)才用 TreeMap。
- TreeMap 非线程安全且迭代器 fail-fast;并发下需要有序 Map 时,JUC 给的答案是跳表实现的 ConcurrentSkipListMap,而不是给 TreeMap 加锁。