注:源码系列文章主要是对某付费专栏的总结记录。如有侵权,请联系删除。
1 知识储备
在了解 TreeMap 之前,我们来看看日常工作中排序的两种方式,作为我们学习的基础储备,两种方式的代码如下:
@Data
class Entry implements Comparable<Entry> {
private final Integer id;
Entry(Integer id) {
this.id = id;
}
@Override
public int compareTo(Entry o) {
// 默认从小到大排序
return id - o.id;
}
}
@Test
public void execute() {
// 第一种排序,从小到大排序,实现 Comparable 的 compareTo 方法进行排序
List<Entry> list1 = new ArrayList<>();
for (int i = 5; i > 0; i--) {
list1.add(new Entry(i));
}
Collections.sort(list1);
System.out.println("list1: " + JSON.toJSONString(list1));
// 第二种排序,从大到小排序,利用外部排序器 Comparator 进行排序
List<Entry> list2 = new ArrayList<>();
for (int i = 5; i > 0; i--) {
list2.add(new Entry(i));
}
Collections.sort(list2, new Comparator<Entry>() {
@Override
public int compare(Entry o1, Entry o2) {
return o2.id - o1.id;
}
});
System.out.println("list2: " + JSON.toJSONString(list2));
}
执行结果:
list1: [{"id":1},{"id":2},{"id":3},{"id":4},{"id":5}]
list2: [{"id":5},{"id":4},{"id":3},{"id":2},{"id":1}]
以上两种就是分别通过 Comparable 和 Comparator 两者进行排序的方式,而 TreeMap 利用的也是此原理,从而实现了对 key 的排序。
2 TreeMap 整理架构
TreeMap 底层的数据结构就是红黑树,和 HashMap 的红黑树结构一样。
不同的是,TreeMap 可利用了红黑树左节点小,右节点大的性能,根据 key 进行排序,使每个元素都能够插入到红黑树大小适当的位置,维护了 key 的大小关系,适用于 key 需要排序的场景。
因为底层使用的是平衡红黑树的结构,所以 containsKey、get、put、remove 等方法的时间复杂度都是 log(n)。
2.1 属性
TreeMap 常见的属性有:
// 比较器。
// 如果外部有传进来 Comparator 比较器,则用外部的
// 如果外部比较器为空,则使用 key 自己实现的 Comparable#compareTof 方法
private final Comparator<? super K> comparator;
// 红黑树的根节点
private transient Entry<K,V> root;
// 红黑树已有元素大小
private transient int size = 0;
// 树结构变化的版本号
private transient int modCount = 0;
// Red-black mechanics
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;
boolean color = BLACK;
Entry(K key, V value, Entry<K,V> parent) {
this.key = key;
this.value = value;
this.parent = parent;
}
}
2.2 新增节点
TreeMap 新增节点源码:
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) {
// 自旋找到 key 应该新增的位置,就是应该挂载到哪个节点的子节点上
do {
// 一次循环结束时,parent 就是上次比过的对象
parent = t;
// 通过 compare 来比较 key 和当前对比节点 t 的大小
cmp = cpr.compare(key, t.key);
// key 小于 t,把 t 左边的节点赋值给 t,因为红黑树左边的值比较小,循环再比
if (cmp < 0)
t = t.left;
// key 大于 t,把 t 右边的节点赋值给 t,因为红黑树右边的值比较小,循环再比
else if (cmp > 0)
t = t.right;
// 如果相等则表示当前节点 t 和要新增的节点相同,直接赋值当前节点 t 的值即可
else
return t.setValue(value);
// t 为空说明已经到叶子节点了
} 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);
}
// 到这一步时已经找到了当前要新增节点 key 的父节点 parent 了
Entry<K,V> e = new Entry<>(key, value, parent);
// cmp 代表最后一次对比的大小,小于 0,代表 e 在上一个节点的左边
if (cmp < 0)
parent.left = e;
// cmp 代表左后一次对比的大小,大于 0,代表 e 在上一个节点的右边
else
parent.right = e;
// 相等的情况上面寻找 parent 时已经做过了
// 着色旋转,达到平衡
fixAfterInsertion(e);
size++;
modCount++;
return null;
}
从源码中,我们可以看到:
- 新增节点时,就是利用了红黑树左小右大的特性,从根节点不断往下查找,直到找到节点是 null 为止,节点为 null 说明到达了叶子节点;
- 查找的过程中,发现 key 值已经存在,直接覆盖值;
- TreeMap 是禁止 key 是 null 值的。
类似的,TreeMap 查找也是类似的原理。
2.3 小结
TreeMap 相对来说比较简单,红黑树和 HashMap 比较类似,比较关键的是通过 compare 来比较 key 的大小,然后利用红黑树左小右大的特性,为每个 key 找到自己的位置,从而维护了 key 的大小排序顺序。
------------------------------------- END -------------------------------------