JDK 8 TreeMap 源码详解(完整版带详细注释)
1. 基本结构和常量定义
public class TreeMap<K,V>
extends AbstractMap<K,V>
implements NavigableMap<K,V>, Cloneable, java.io.Serializable {
// 序列化版本号
private static final long serialVersionUID = 919286545866124006L;
// 比较器,用于定义键的排序规则
// 如果为null,则使用键的自然排序(要求键实现Comparable接口)
private final Comparator<? super K> comparator;
// 红黑树的根节点
private transient Entry<K,V> root;
// 红黑树中节点的数量
private transient int size = 0;
// 修改次数,用于快速失败机制
private transient int modCount = 0;
// 红黑树节点颜色常量
private static final boolean RED = false;
private static final boolean BLACK = true;
}
2. Entry节点类
/**
* 红黑树节点
* @param <K> 键的类型
* @param <V> 值的类型
*/
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; // 节点颜色,默认为黑色
/**
* 构造函数
*/
Entry(K key, V value, Entry<K,V> parent) {
this.key = key;
this.value = value;
this.parent = parent;
}
/**
* 获取键
*/
public K getKey() {
return key;
}
/**
* 获取值
*/
public V getValue() {
return value;
}
/**
* 设置值并返回旧值
*/
public V setValue(V value) {
V oldValue = this.value;
this.value = value;
return oldValue;
}
/**
* 判断两个节点是否相等
*/
public boolean equals(Object o) {
if (!(o instanceof Map.Entry))
return false;
Map.Entry<?,?> e = (Map.Entry<?,?>)o;
return valEquals(key,e.getKey()) && valEquals(value,e.getValue());
}
/**
* 计算哈希码
*/
public int hashCode() {
return (key==null ? 0 : key.hashCode()) ^
(value==null ? 0 : value.hashCode());
}
/**
* 字符串表示
*/
public String toString() {
return key + "=" + value;
}
}
3. 构造函数
/**
* 无参构造函数
* 使用键的自然排序
*/
public TreeMap() {
comparator = null;
}
/**
* 指定比较器的构造函数
* @param comparator 比较器
*/
public TreeMap(Comparator<? super K> comparator) {
this.comparator = comparator;
}
/**
* 从Map构造TreeMap
* @param m 要构造TreeMap的Map
*/
public TreeMap(Map<? extends K, ? extends V> m) {
comparator = null;
putAll(m);
}
/**
* 从SortedMap构造TreeMap
* @param m 要构造TreeMap的SortedMap
*/
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) {
}
}
4. 核心方法实现
4.1 比较方法
/**
* 使用比较器或自然排序比较两个键
* @param k1 键1
* @param k2 键2
* @return 比较结果
*/
@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);
}
/**
* 判断两个值是否相等
*/
static final boolean valEquals(Object o1, Object o2) {
return (o1==null ? o2==null : o1.equals(o2));
}
4.2 get方法
/**
* 获取指定键对应的值
* @param key 键
* @return 对应的值,如果不存在返回null
*/
public V get(Object key) {
Entry<K,V> p = getEntry(key);
return (p==null ? null : p.value);
}
/**
* 根据键获取对应的节点
* @param key 键
* @return 对应的节点,如果不存在返回null
*/
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;
}
/**
* 使用比较器获取节点
*/
@SuppressWarnings("unchecked")
final Entry<K,V> getEntryUsingComparator(Object key) {
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;
}
4.3 put方法
/**
* 添加键值对到TreeMap中
* @param key 键
* @param value 值
* @return 旧值,如果不存在返回null
*/
public V put(K key, V value) {
Entry<K,V> t = root;
// 如果树为空,创建根节点
if (t == null) {
compare(key, key); // type (and possibly null) check
root = new Entry<>(key, value, null);
size = 1;
modCount++;
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
return t.setValue(value);
} 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);
}
// 创建新节点
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;
}
4.4 remove方法
/**
* 移除指定键的映射关系
* @param key 要移除的键
* @return 被移除的值,如果不存在返回null
*/
public V remove(Object key) {
Entry<K,V> p = getEntry(key);
if (p == null)
return null;
V oldValue = p.value;
deleteEntry(p);
return oldValue;
}
/**
* 删除指定节点
* @param p 要删除的节点
*/
private void deleteEntry(Entry<K,V> p) {
modCount++;
size--;
// 如果节点有两个子节点,找到后继节点替换
if (p.left != null && p.right != null) {
Entry<K,V> s = successor(p);
p.key = s.key;
p.value = s.value;
p = s;
} // p has 2 children
// 开始删除节点p
Entry<K,V> replacement = (p.left != null ? p.left : p.right);
if (replacement != null) {
// 如果p有一个子节点
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;
// 清空p的引用,帮助GC
p.left = p.right = p.parent = null;
// 如果p是黑色节点,需要修复红黑树性质
if (p.color == BLACK)
fixAfterDeletion(replacement);
} else if (p.parent == null) { // return if we are the only node.
// 如果p是根节点且没有子节点
root = null;
} else { // No children. Use self as phantom replacement and unlink.
// 如果p是叶子节点
if (p.color == BLACK)
fixAfterDeletion(p);
// 从父节点中移除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;
}
}
}
5. 红黑树操作
5.1 插入后修复
/**
* 插入节点后修复红黑树性质
* @param x 新插入的节点
*/
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) {
// 情况1:叔叔节点是红色
setColor(parentOf(x), BLACK);
setColor(y, BLACK);
setColor(parentOf(parentOf(x)), RED);
x = parentOf(parentOf(x));
} else {
// 情况2:叔叔节点是黑色
if (x == rightOf(parentOf(x))) {
x = parentOf(x);
rotateLeft(x);
}
setColor(parentOf(x), BLACK);
setColor(parentOf(parentOf(x)), RED);
rotateRight(parentOf(parentOf(x)));
}
} else {
// 父节点是祖父节点的右子节点(镜像情况)
Entry<K,V> y = leftOf(parentOf(parentOf(x)));
if (colorOf(y) == RED) {
setColor(parentOf(x), BLACK);
setColor(y, BLACK);
setColor(parentOf(parentOf(x)), RED);
x = parentOf(parentOf(x));
} else {
if (x == leftOf(parentOf(x))) {
x = parentOf(x);
rotateRight(x);
}
setColor(parentOf(x), BLACK);
setColor(parentOf(parentOf(x)), RED);
rotateLeft(parentOf(parentOf(x)));
}
}
}
root.color = BLACK; // 根节点始终是黑色
}
5.2 删除后修复
/**
* 删除节点后修复红黑树性质
* @param x 要修复的节点
*/
private void fixAfterDeletion(Entry<K,V> x) {
while (x != root && colorOf(x) == BLACK) {
if (x == leftOf(parentOf(x))) {
// x是父节点的左子节点
Entry<K,V> sib = rightOf(parentOf(x));
if (colorOf(sib) == RED) {
// 情况1:兄弟节点是红色
setColor(sib, BLACK);
setColor(parentOf(x), RED);
rotateLeft(parentOf(x));
sib = rightOf(parentOf(x));
}
if (colorOf(leftOf(sib)) == BLACK &&
colorOf(rightOf(sib)) == BLACK) {
// 情况2:兄弟节点的两个子节点都是黑色
setColor(sib, RED);
x = parentOf(x);
} else {
if (colorOf(rightOf(sib)) == BLACK) {
// 情况3:兄弟节点的右子节点是黑色
setColor(leftOf(sib), BLACK);
setColor(sib, RED);
rotateRight(sib);
sib = rightOf(parentOf(x));
}
// 情况4:兄弟节点的右子节点是红色
setColor(sib, colorOf(parentOf(x)));
setColor(parentOf(x), BLACK);
setColor(rightOf(sib), BLACK);
rotateLeft(parentOf(x));
x = root;
}
} else { // symmetric
// x是父节点的右子节点(镜像情况)
Entry<K,V> sib = leftOf(parentOf(x));
if (colorOf(sib) == RED) {
setColor(sib, BLACK);
setColor(parentOf(x), RED);
rotateRight(parentOf(x));
sib = leftOf(parentOf(x));
}
if (colorOf(rightOf(sib)) == BLACK &&
colorOf(leftOf(sib)) == BLACK) {
setColor(sib, RED);
x = parentOf(x);
} else {
if (colorOf(leftOf(sib)) == BLACK) {
setColor(rightOf(sib), BLACK);
setColor(sib, RED);
rotateLeft(sib);
sib = leftOf(parentOf(x));
}
setColor(sib, colorOf(parentOf(x)));
setColor(parentOf(x), BLACK);
setColor(leftOf(sib), BLACK);
rotateRight(parentOf(x));
x = root;
}
}
}
setColor(x, BLACK);
}
5.3 旋转操作
/**
* 左旋操作
* @param p 要旋转的节点
*/
private void rotateLeft(Entry<K,V> p) {
if (p != null) {
Entry<K,V> r = p.right;
p.right = r.left;
if (r.left != null)
r.left.parent = p;
r.parent = p.parent;
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.parent = r;
}
}
/**
* 右旋操作
* @param p 要旋转的节点
*/
private void rotateRight(Entry<K,V> p) {
if (p != null) {
Entry<K,V> l = p.left;
p.left = l.right;
if (l.right != null) l.right.parent = p;
l.parent = p.parent;
if (p.parent == null)
root = l;
else if (p.parent.right == p)
p.parent.right = l;
else p.parent.left = l;
l.right = p;
p.parent = l;
}
}
6. 辅助方法
6.1 基本操作辅助方法
/**
* 获取节点的父节点
*/
private static <K,V> Entry<K,V> parentOf(Entry<K,V> p) {
return (p == null ? null: p.parent);
}
/**
* 设置节点颜色
*/
private static <K,V> void setColor(Entry<K,V> p, boolean c) {
if (p != null)
p.color = c;
}
/**
* 获取节点颜色
*/
private static <K,V> boolean colorOf(Entry<K,V> p) {
return (p == null ? BLACK : p.color);
}
/**
* 获取节点的左子节点
*/
private static <K,V> Entry<K,V> leftOf(Entry<K,V> p) {
return (p == null) ? null: p.left;
}
/**
* 获取节点的右子节点
*/
private static <K,V> Entry<K,V> rightOf(Entry<K,V> p) {
return (p == null) ? null: p.right;
}
6.2 前驱和后继节点
/**
* 获取指定节点的前驱节点(中序遍历的前一个节点)
* @param t 节点
* @return 前驱节点
*/
static <K,V> Entry<K,V> predecessor(Entry<K,V> t) {
if (t == null)
return null;
else if (t.left != null) {
// 如果有左子树,前驱是左子树中的最大节点
Entry<K,V> p = t.left;
while (p.right != null)
p = p.right;
return p;
} else {
// 如果没有左子树,前驱是某个祖先节点
Entry<K,V> p = t.parent;
Entry<K,V> ch = t;
while (p != null && ch == p.left) {
ch = p;
p = p.parent;
}
return p;
}
}
/**
* 获取指定节点的后继节点(中序遍历的后一个节点)
* @param t 节点
* @return 后继节点
*/
static <K,V> 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;
}
}
6.3 最小和最大节点
/**
* 获取树中的最小节点
* @param t 树的根节点
* @return 最小节点
*/
static <K,V> Entry<K,V> getFirstEntry(Entry<K,V> t) {
if (t != null)
while (t.left != null)
t = t.left;
return t;
}
/**
* 获取树中的最大节点
* @param t 树的根节点
* @return 最大节点
*/
static <K,V> Entry<K,V> getLastEntry(Entry<K,V> t) {
if (t != null)
while (t.right != null)
t = t.right;
return t;
}
7. NavigableMap接口方法实现
7.1 范围查询相关方法
/**
* 获取第一个键值对
*/
public Map.Entry<K,V> firstEntry() {
return exportEntry(getFirstEntry(root));
}
/**
* 获取最后一个键值对
*/
public Map.Entry<K,V> lastEntry() {
return exportEntry(getLastEntry(root));
}
/**
* 获取并移除第一个键值对
*/
public Map.Entry<K,V> pollFirstEntry() {
Entry<K,V> p = getFirstEntry(root);
Map.Entry<K,V> result = exportEntry(p);
if (p != null)
deleteEntry(p);
return result;
}
/**
* 获取并移除最后一个键值对
*/
public Map.Entry<K,V> pollLastEntry() {
Entry<K,V> p = getLastEntry(root);
Map.Entry<K,V> result = exportEntry(p);
if (p != null)
deleteEntry(p);
return result;
}
/**
* 导出Entry为不可变的SimpleImmutableEntry
*/
static <K,V> Map.Entry<K,V> exportEntry(TreeMap.Entry<K,V> e) {
return (e == null) ? null :
new AbstractMap.SimpleImmutableEntry<>(e);
}
7.2 排序相关方法
/**
* 获取比较器
*/
public Comparator<? super K> comparator() {
return comparator;
}
/**
* 获取第一个键
*/
public K firstKey() {
return key(getFirstEntry(root));
}
/**
* 获取最后一个键
*/
public K lastKey() {
return key(getLastEntry(root));
}
/**
* 获取键(null安全)
*/
static <K> K key(TreeMap.Entry<K,?> e) {
if (e == null)
throw new NoSuchElementException();
return e.key;
}
8. 迭代器实现
/**
* 键的迭代器
*/
Iterator<K> keyIterator() {
return new KeyIterator(getFirstEntry(root));
}
/**
* 值的迭代器
*/
Iterator<V> valueIterator() {
return new ValueIterator(getFirstEntry(root));
}
/**
* 键值对的迭代器
*/
Iterator<Map.Entry<K,V>> entryIterator() {
return new EntryIterator(getFirstEntry(root));
}
/**
* 键迭代器实现
*/
final class KeyIterator extends PrivateEntryIterator<K> {
KeyIterator(Entry<K,V> first) {
super(first);
}
public K next() {
return nextEntry().key;
}
}
/**
* 值迭代器实现
*/
final class ValueIterator extends PrivateEntryIterator<V> {
ValueIterator(Entry<K,V> first) {
super(first);
}
public V next() {
return nextEntry().value;
}
}
/**
* 键值对迭代器实现
*/
final class EntryIterator extends PrivateEntryIterator<Map.Entry<K,V>> {
EntryIterator(Entry<K,V> first) {
super(first);
}
public Map.Entry<K,V> next() {
return nextEntry();
}
}
/**
* 私有迭代器基类
*/
abstract class PrivateEntryIterator<T> implements Iterator<T> {
Entry<K,V> next;
Entry<K,V> lastReturned;
int expectedModCount;
PrivateEntryIterator(Entry<K,V> first) {
expectedModCount = modCount;
lastReturned = null;
next = first;
}
public final boolean hasNext() {
return next != null;
}
final Entry<K,V> nextEntry() {
Entry<K,V> e = next;
if (e == null)
throw new NoSuchElementException();
if (modCount != expectedModCount)
throw new ConcurrentModificationException();
next = successor(e);
lastReturned = e;
return e;
}
public void remove() {
if (lastReturned == null)
throw new IllegalStateException();
if (modCount != expectedModCount)
throw new ConcurrentModificationException();
// deleted entries are replaced by their successors
if (lastReturned.left != null && lastReturned.right != null)
next = lastReturned;
deleteEntry(lastReturned);
expectedModCount = modCount;
lastReturned = null;
}
}
9. 其他重要方法
9.1 size和isEmpty方法
/**
* 返回TreeMap中键值对的数量
*/
public int size() {
return size;
}
/**
* 判断TreeMap是否为空
*/
public boolean isEmpty() {
return size == 0;
}
9.2 clear方法
/**
* 清空TreeMap中的所有元素
*/
public void clear() {
modCount++;
size = 0;
root = null;
}
9.3 clone方法
/**
* 克隆方法
*/
public Object clone() {
TreeMap<?,?> clone;
try {
clone = (TreeMap<?,?>) super.clone();
} catch (CloneNotSupportedException e) {
throw new InternalError(e);
}
// Put clone into "virgin" state (except for comparator)
clone.root = null;
clone.size = 0;
clone.modCount = 0;
clone.entrySet = null;
clone.navigableKeySet = null;
clone.descendingMap = null;
// Initialize clone with our mappings
try {
clone.buildFromSorted(size, entrySet().iterator(), null, null);
} catch (java.io.IOException cannotHappen) {
} catch (ClassNotFoundException cannotHappen) {
}
return clone;
}
10. 总结
10.1 TreeMap的特点
- 底层实现:基于红黑树实现
- 有序性:保持键的排序顺序(自然排序或自定义比较器)
- 唯一性:不允许存储重复的键
- 不允许null键:如果使用自然排序,键不能为null
- 允许null值:可以存储null值
- 非线程安全:不是线程安全的
- 高性能:基本操作(get、put、remove)时间复杂度为O(log n)
10.2 时间复杂度分析
- 添加元素:O(log n)
- 删除元素:O(log n)
- 查找元素:O(log n)
- 遍历元素:O(n)
- 范围查询:O(log n + k),k为结果集大小
10.3 空间复杂度
- 存储空间:O(n) - n为存储的键值对个数
- 额外空间:每个节点需要额外的parent、left、right指针和颜色信息
10.4 使用建议
-
适用场景:
- 需要按键排序存储数据
- 需要频繁进行范围查询
- 需要找到最小/最大元素
- 需要实现有序映射
-
不适用场景:
- 不需要排序功能(考虑使用HashMap)
- 需要快速随机访问(考虑使用ArrayList)
- 多线程环境(考虑使用Collections.synchronizedSortedMap()包装)
-
性能优化:
- 确保键的compareTo()或比较器实现正确
- 对于大量数据的批量插入,可以考虑使用buildFromSorted()方法
- 选择合适的键类型以获得良好的排序性能
10.5 红黑树的性质
- 节点颜色:每个节点要么是红色,要么是黑色
- 根节点:根节点是黑色
- 叶子节点:所有叶子节点(NIL节点)都是黑色
- 红色节点:红色节点的两个子节点都是黑色
- 黑色高度:从任一节点到其每个叶子的所有路径都包含相同数目的黑色节点
TreeMap通过这些性质保证了树的平衡性,从而确保了各种操作的对数时间复杂度。它是Java集合框架中实现有序映射的经典数据结构。