TreeMap源码阅读

前言

TreeMap算是要分析的Map家族中的最后一个成员了,它是AbstractMap的子类,并实现了NavigableMap接口,是Map家族中比较特殊的,支持元素排序插入的Map集合类。博客原文地址TreeMap源码阅读

TreeMap简介

TreeMap的类图如下:

  • 虚线箭头代表的是实现接口;实线箭头代表的是继承关系

由上图所示,它是AbstractMap的子类,因此就具备了一个普通Map的全部操作,如put、get、containsKey等方法,同时,它也实现了NavigableMap,由图中关系,NavigableMap是SortedMap的子接口,所以NavigableMap具备排序功能,且有它的名字可知,它也具备点查询(Navigable)功能。下面简单介绍一下NavigableMap。

NavigableMap在SortedMap的基础上,增加了一些navigation方法,比如说lowerEntry、floorEntry、ceilingEntry和higherEntry,它们传入的参数都是一个键,类型为泛型K,返回值都为一个Map.Entry对象。这四个方法分别返回小于、小于或等于、大于或等于、大于传入的key的最接近的Entry,也就是说floorEntry方法,如果有小于key的entry1和等于key的entry2,则返回entry2。

源码分析

根据名字也可猜出,TreeMap的存储结构为树,树的节点为一个内部类Entry,它是Map.Entry的一个子类。它的属性包含了key,value,指向左兄弟节点的left和指向右兄弟节点的right,指向父节点的parent,表示颜色的布尔值color(为true表示黑色,false表示红色),因此可知TreeMap是一个红黑树。在HashMap的分析中,1.8之后也增加了红黑树的实现,在此可以进行一下对比。TreeMap中没有指向前一个节点的prev指针,这比较容易理解,HashMap中红黑树是对开链法形成的大于8个节点的链表所进行的,其原型还是链表,只是通过增加了指针形成了红黑树,当节点小于8时,还需要进行untreeify,因此指向前一个节点的指针存储的为链表信息。而TreeMap本身就是树状结构,因此不需要保存链表信息的prev属性。另一个区别是TreeMap中表示颜色的属性为color,虽然也为布尔值,但它的赋值是通过定义好的两个常量进行的,常量定义如下

1
2
private static final boolean RED = false;
private static final boolean BLACK = true;

修改颜色时将对应的常量赋值给color属性即可,这样的定义方式更加形象、直观。至于区别的原因,暂时理解为不同contributor的编程习惯。TreeMap的内部类Entry实现非常简单,只是实现了getValue、getKey、setValue、equals和hashCode方法。其中equals方法首先要求传入的为Map.Entry子对象,其次要求key和value同时相等时才返回true。hashCode返回的为key,value的hashCode的异或值。由于Entry的实现简洁,树的结构操作都放在了TreeMap的方法中,因此不禁产生疑问,TreeMap是如何进行插入和查询的呢?

带着疑问先来看get方法,其核心为getEntry方法,代码如下

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
final Entry<K,V> getEntry(Object key) {
// Offload comparator-based version for sake of performance
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;
}

getEntry实现的是二叉搜索树的查找算法,也就是与当前节点的key值比较,如果小于,则在左子树中遍历查找,大于的话则在右子树查找,等于返回当前节点。这里需要注意的是,为了查找的性能,这里有一个判断语句,如果comparator不为空,则调用了getEntryUsingComparator方法,该方法的实现如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
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)
p = p.left;
else if (cmp > 0)
p = p.right;
else
return p;
}
}
return null;
}

可以看到,在算法上,此方法仍然使用的是二叉搜索树的查找算法,因此在性能上不会有提升,它的性能提高主要体现在comparator的比较算法上,但大多数情况下comparator的比较算法不会有太大差别,本方法的存在只是为了以防万一。由此,get方法实现的是二叉搜索树的标准查找算法,下面分析put方法。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
public V put(K key, V value) {
Entry<K,V> t = root;
//TreeMap为空
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;
//使用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);
}
//key为Comparable,使用key.compareTo比较
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;
}

put方法的算法比价清晰,首先查找key值,如果找到了则用新值替换旧值,否则,在合适的位置插入一个新的Entry。由于TreeMap为红黑树,因此需要在插入节点之后进行fix操作,以使得插入后的仍然满足红黑树的特性,该方法为fixAfterInsertion,其定义如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
/** From CLR */
private void fixAfterInsertion(Entry<K,V> x) {
//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) {
setColor(parentOf(x), BLACK);
setColor(y, BLACK);
setColor(parentOf(parentOf(x)), RED);
x = parentOf(parentOf(x));
} else {
//x为右子节点,x.parent为左子节点,需要两次旋转进行调整
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;
}

上面代码加上了必要的注释,总的来说就是分为下面几步解决红黑树的平衡问题:

  1. 新插入的节点设为红色
  2. 通过循环,从新插入节点向上遍历,解决冲突,循环的终止条件是x.parent.color = BLACK,这里有一个隐含条件就是x.color = RED,因此当x的父节点为黑色时,红黑树已经调整完毕,其他的跳出循环为边界情况
  3. 循环中,根据x父节点是否为其祖父节点的左子节点分成了两个分支,在各自的分支中,根据x叔父节点是否为红色,又分为两个分支,为红色时不需要进行旋转,当前结构平衡,只需要进行颜色的修改;不为红色时,需要通过旋转来保持红黑树的平衡。这里的算法可以参考博客红黑树数据结构剖析

这里可以对照我的另一篇关于HashMap的博客HashMap解析,其中也有红黑树的相关实现,其算法一致,只是TreeMap中的更加直观,也更加具有语义化,比较容易理解。

get、put之后,我们来分析一下查询方法,get方法实际也有查询的作用,但由于containsValue方法并没有使用get方法,因此单独分析一下,这里也比较容易理解,因为TreeMap的节点是通过key的大小来进行排序的,因此查询value的存在时,不能通过二叉搜索算法来查询,只能通过遍历的方式,时间复杂度为Ω(n)。containsValue定义如下:

1
2
3
4
5
6
public boolean containsValue(Object value) {
for (Entry<K,V> e = getFirstEntry(); e != null; e = successor(e))
if (valEquals(value, e.value))
return true;
return false;
}

  • getFirstEntry方法返回树的最左边的叶子节点,也就是key最小的节点
  • successor(e)返回e的下一个节点,也就是key值紧挨着e.key,且大于e.key的节点

算法首先找到树中最左边叶子节点,然后successor每次返回当前节点的下一个节点,遍历整个树,直到找到该值或者到达树的最后一个节点。由于successor方法的时间复杂度并不是O(1)的,因此containsValue的时间复杂度大于n,又由于红黑树是二叉平衡树,因此successor时间复杂度不会到达O(n)的级别,因此containsValue的时间复杂度为Ω(n)。successor方法的定义如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
static <K,V> TreeMap.Entry<K,V> successor(Entry<K,V> t) {
if (t == null)
return null;
//如果t有右子树,则返回其右子树中的最左边的叶子节点
else if (t.right != null) {
Entry<K,V> p = t.right;
while (p.left != null)
p = p.left;
return p;
} else { //t不存在右子树,需要向上遍历,找到t为左子节点的第一个父节点
Entry<K,V> p = t.parent;
Entry<K,V> ch = t;
while (p != null && ch == p.right) {
ch = p;
p = p.parent;
}
return p;
}
}

算法比较简单,已经在注释中说明,在此不再赘述。

下面分析remove方法,由于remove方法带来了红黑树的结构变化,因此必然也涉及到了删除节点之后的fix过程。remove方法的核心部分为deleteEntry方法,其定义如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
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;
p.value = s.value;
p = s;
} // p has 2 children,实际删除的对象变为successor(p)
// Start fixup at replacement node, if it exists.
Entry<K,V> replacement = (p.left != null ? p.left : p.right);
if (replacement != null) { //p有孩子节点
// 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,p为黑色的时候,删除p会影响到路径上黑色节点的个数,需要进行fix操作
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;
}
}
}

上面代码中,要实现的是删除节点p,在删除的过程中,p的状态有三种情况:p有两个子节点、有一个子节点、没有子节点。

  1. p有两个子节点的时候,会找到p的右子树中最左边的节点n,将n的值覆盖掉p的值,此时实际删除的节点变成了n,n没有左子节点,因此将n的右子树的根节点nLeftRoot嫁接到n的位置即可。
  2. p有一个子节点的时候,直接删除的节点为p,将p的左/右子树的根节点childRoot嫁接到p的位置
  3. p没有子节点的时候,需要先进行调整,然后删除节点p

红黑树调整也不是每次删除都会发生的,只有当删除的节点p为黑色的时候,会导致其所在路径上的黑色节点数少一个,此时需要进行调整,通过方法fixAfterDeletion来进行,其定义如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
/** From CLR */
private void fixAfterDeletion(Entry<K,V> x) { //x所在路径黑色节点少一个
while (x != root && colorOf(x) == BLACK) { //x为红色,直接终止循环,将x设为黑色,调整完毕
//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));
}
//此时sib必为黑色,若sib的两个子节点都为黑色,则此路径上黑色节点数多1个
if (colorOf(leftOf(sib)) == BLACK &&
colorOf(rightOf(sib)) == BLACK) {
setColor(sib, RED); //将sib设为红色,黑色节点数-1
x = parentOf(x); //x为sib的兄弟节点,x设为x.parent,递归调整红黑树
} else { //sib的子节点一红一黑
if (colorOf(rightOf(sib)) == BLACK) { //若右子节点为黑色
setColor(leftOf(sib), BLACK); //设左子节点为黑色, sib为红色
setColor(sib, RED);
rotateRight(sib); //右旋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
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);
}

该方法传入节点x,由删除方法的触发条件可知,本方法需要解决的问题是,由于x节点的删除,该节点所在的路径上黑色节点数少了一个,因此需要恢复从根节点到x节点上的黑色节点数,使其加一,这样才能恢复树的平衡状态。算法通过一个while循环来进行调整,终止条件为x到达根节点或者x的颜色为红色。x到达根节点时,直接将其设为黑色,方法返回,这比较好理解,因为到达根节点时说明下面各层的颜色调整结束,整棵树满足红黑树。而x为红色时,为什么也可以跳出呢?因为要调整的是x所在路径上的黑色节点数+1,因此如果x为红色的话,只需要将x设置为黑色,则其所在路径上黑色节点数达到要求,因此整棵树也达到了平衡状态。

算法的解释通过注释在上面进行了必要的解释,具体算法可以参考红黑树数据结构剖析

总结

TreeMap中很多东西和HashMap中的一部分东西非常类似,这是因为HashMap在开链法的链表长度大于8的时候,会将链表树化为一个红黑树,因此TreeMap中的插入,删除相当于HashMap中树化的插入、删除。但它们两者也大有不同,一个区别就是TreeMap是一个排序的map,可以通过指定comparator来指定排序规则,而HashMap中树化的比较是通过key的hashCode来进行的,且HashMap首先是一个Hash表,其次bucket对应的链表才是一个红黑树,因此HashMap可以理解为O(1)时间内查找到key值,而TreeMap是一个排序的Map,如果需要维护一个排好序的map时,可以使用TreeMap,且由于红黑树的性质,它可以保证O(logn)的时间内查找到你需要定位的Entry(通过key值来进行比较)。

live long and prosper