[简单集合] TreeMap源码分析
1 存储结构
红黑树
TreeMap只使用到了红黑树，所以时间复杂度为O(log n)。红黑树的特性为：
- 每个节点是黑色或者红色
- 根节点是黑色
- 每个叶子节点（NIL）是黑色（注意：这里叶子节点，是指为空的叶子节点）
- 如果一个节点是红色的，则它的子节点必须是黑色的
- 从一个节点到该节点的子孙节点的所有路径上包含相同数目的黑节点
2 源码解析
2.1 属性
// 比较器,如果没传则key要实现Comparable接口
private final Comparator<? super K> comparator;
// 根节点
private transient Entry<K,V> root;
// 元素个数
private transient int size = 0;
// 修改次数
private transient int modCount = 0;
（1）comparator
按key的大小排序有两种方式，一种是key实现Comparable接口，一种方式通过构造方法传入比较器。
（2）root
根节点，TreeMap没有桶的概念，所有的元素都存储在一颗树中。
2.2 Entry内部类
典型的红黑树结构
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;
}
2.3 构造方法
// 默认构造方法,key必须实现Comparable接口
public TreeMap() {
comparator = null;
}
// 使用传入的comparator比较两个key的大小
public TreeMap(Comparator<? super K> comparator) {
this.comparator = comparator;
}
// key必须实现Comparable接口，把传入map中的所有元素保存到新的TreeMap中
public TreeMap(Map<? extends K, ? extends V> m) {
comparator = null;
putAll(m);
}
// 使用传入map的比较器，并把传入map中的所有元素保存到新的TreeMap中
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) {
}
}
构造方法主要分成两类，一类是使用comparator比较器，一类是key必须实现Comparable接口。
2.4 get(Object key)方法
获取元素，典型的二叉查找树的查找方法。
public V get(Object key) {
// 根据key查找元素
Entry<K,V> p = getEntry(key);
// 找到了返回value值,没找到返回null
return (p==null ? null : p.value);
}
final Entry<K,V> getEntry(Object key) {
// Offload comparator-based version for sake of performance
// 如果比较器不为空,使用带有comparator的版本获取元素
if (comparator != null)
return getEntryUsingComparator(key);
if (key == null)
throw new NullPointerException();
@SuppressWarnings("unchecked")
Comparable<? super K> k = (Comparable<? super K>) key; // 将key强转为Comparable类型,如果没有传入comparator,key必须实现Comparable接口
// 从根元素开始遍历
Entry<K,V> p = root;
while (p != null) {
int cmp = k.compareTo(p.key);
if (cmp < 0)
// 如果小于0从左子树查找
p = p.left;
else if (cmp > 0)
// 如果大于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)
// 如果小于0从左子树查找
p = p.left;
else if (cmp > 0)
// 如果大于0从右子树查找
p = p.right;
else
return p;
}
}
return null;
}
（1）从root遍历整个树；
（2）如果待查找的key比当前遍历的key小，则在其左子树中查找；
（3）如果待查找的key比当前遍历的key大，则在其右子树中查找；
（4）如果待查找的key与当前遍历的key相等，则找到了该元素，直接返回；
（5）从这里可以看出是否有comparator分化成了两个方法，但是内部逻辑一模一样，因此可见笔者 comparator=(k1,k2)->((Comparable<?superK>)k1).compareTo(k2);这种改造的必要性。
2.5 put(K key,V value)方法
插入元素流程：
插入后再平衡流程：
插入后再平衡
（如果父节点是祖父节点的左节点）
| 情况 | 策略 |
|---|---|
| 1）父节点为红色，叔叔节点也为红色 | （1）将父节点设为黑色；（2）将叔叔节点设为黑色；（3）将祖父节点设为红色；（4）将祖父节点设为新的当前节点，进入下一次循环判断； |
| 2）父节点为红色，叔叔节点为黑色，且当前节点是其父节点的右节点 | （1）将父节点作为新的当前节点（2）以新当节点为支点进行左旋，进入情况3）； |
| 3）父节点为红色，叔叔节点为黑色，且当前节点是其父节点的左节点 | （1）将父节点设为黑色；（2）将祖父节点设为红色；（3）以祖父节点为支点进行右旋，进入下一次循环判断； |
（如果父节点是祖父节点的右节点，则正好与上面反过来）
| 情况 | 策略 |
|---|---|
| 1）父节点为红色，叔叔节点也为红色 | （1）将父节点设为黑色；（2）将叔叔节点设为黑色；（3）将祖父节点设为红色；（4）将祖父节点设为新的当前节点，进入下一次循环判断； |
| 2）父节点为红色，叔叔节点为黑色，且当前节点是其父节点的左节点 | （1）将父节点作为新的当前节点（2）以新当节点为支点进行右旋，进入情况3）； |
| 3）父节点为红色，叔叔节点为黑色，且当前节点是其父节点的右节点 | （1）将父节点设为黑色；（2）将祖父节点设为红色；（3）以祖父节点为支点进行左旋，进入下一次循环判断； |
实例：
我们依次向红黑树中插入 4、2、3 三个元素，来一起看看整个红黑树平衡的过程。
三个元素都插入完成后，符合父节点是祖父节点的左节点，叔叔节点为黑色，且当前节点是其父节点的右节点，即情况2）。
情况2）需要做以下两步处理：
（1）将父节点作为新的当前节点；
（2）以新当节点为支点进行左旋，进入情况3）；
情况3）需要做以下三步处理：
（1）将父节点设为黑色；
（2）将祖父节点设为红色；
（3）以祖父节点为支点进行右旋，进入下一次循环判断；
下一次循环不符合父节点为红色了，退出循环，插入再平衡完成。
