Java 容器源码分析之 TreeMap

TreeMap 是一种基于红黑树实现的 Key-Value 结构。在使用集合视图在 HashMap 中迭代时,是不能保证迭代顺序的; LinkedHashMap 使用了双向链表,保证按照插入顺序或者访问顺序进行迭代。但是有些时候,我们可能需要按照键的大小进行按序迭代,或者在使用哈希表的同时希望按键值进行排序,这个时候 TreeMap 就有其用武之地了。 TreeMap 支持按键值进行升序访问,或者由传入的比较器(Comparator)来控制。

下面基于 JDK 8 的源码对 TreeMap 进行一个简单的分析。

1
2
3
public class TreeMap<K,V>
extends AbstractMap<K,V>
implements NavigableMap<K,V>, Cloneable, java.io.Serializable

同 HashMap 一样, TreeMap 也继承了 AbstractMap,并实现了 Cloneable, Serializable 接口。不同的是, TreeMap 还实现 NavigableMap 接口。

SortedMap 是一个扩展自 Map 的一个接口,对该接口的实现要保证所有的 Key 是完全有序的。

这个顺序一般是指 Key 的自然序(实现 Comparable 接口)或在创建 SortedMap 时指定一个比较器(Comparator)。当我们使用集合的视角(Collection View,由 entrySet、keySet 与 values 方法提供)来迭代时,就可以按序访问其中的元素。

插入 SortedMap 中的所有 Key 的类都必须实现 Comparable 接口(或者可以作为指定的 Comparator 的参数)。在比较两个 Key 时通过调用 k1.compareTo(k2) (or comparator.compare(k1, k2)),因而所有的 Key 都必须能够相互比较,否则会抛出 ClassCastException 的异常。

SortedMap 中 Key 的顺序必须和 equals 保持一致(consistent with equals),
k1.compareTo(k2) == 0 (or comparator.compare(k1, k2)) 和 k1.equals(k2) 要有相同的布尔值。(Comparable 接口的实现不强制要求这一点,但通常都会遵守。)这是因为 Map 接口的定义中,比较 Key 是通过 equals 方法,而在 SortedMap 中比较 Key 则是通过 compareTo (or compare) 方法。如果不一致的,就破坏了 Map 接口的约定。

通过 SortedMap 可以获取其中的一段数据,如 subMap(K fromKey, K toKey), headMap(K toKey), tailMap(K fromKey) 等,所有的区间操作都是左闭右开的。也可以通过 firstKey()lastKey() 来获取第一个和最后一个键。

NavigableMap 是 JDK 1.6 之后新增的接口,扩展了 SortedMap 接口,提供了一些导航方法(navigation methods)来返回最接近搜索目标的匹配结果。

  • lowerEntry(K key) (or lowerKey(K key)),小于给定 Key 的 Entry (or Key)
  • floorEntry(K key) (or floorKey(K key)),小于等于给定 Key 的 Entry (or Key)
  • higherEntry(K key) (or higherKey(K key)),大于给定 Key 的 Entry (or Key)
  • ceilingEntry(K key) (or ceilingKey(K key)),大于等于给定 Key 的 Entry (or Key)

这些方法都有重载的版本,来控制是否包含端点。subMap(K fromKey, K toKey), headMap(K toKey), tailMap(K fromKey) 等方法也是如此。

NavigableMap 可以按照 Key 的升序或降序进行访问和遍历。 descendingMap()descendingKeySet() 则会获取和原来的顺序相反的集合,集合中的元素则是同样的引用,在该视图上的修改会影响到原始的数据。

底层结构

TreeMap 是基于红黑树来实现的,排序时按照键的自然序(要求实现 Comparable 接口)或者提供一个 Comparator 用于排序。

1
2
3
4
5
6
7
8
9
10
11
//比较器,没有指定的话默认使用Key的自然序
private final Comparator<? super K> comparator;

//红黑树根节点
private transient Entry<K,V> root;

//树中节点的数量
private transient int size = 0;

//结构化修改的次数
private transient int modCount = 0;

TreeMap 同样不是线程安全的,基于结构化修改的次数来实现 fail-fast 机制。因而要在多线程环境下使用时,可能需要手动进行同步,或者使用 Collections.synchronizedSortedMap 进行包装。

TreeMap 中的红黑树使用的是「算法导论」中的实现,除了左右链接、红黑标识以外,还有一个指向父节点的连接。红黑树的具体插入及删除细节这里不作过多的解释,更深入的细节可以参考「算法导论」一书,不过建议先看一下 Sedgewick 的讲解。

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
64
65
66
67
68
69
//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;//颜色,指向该节点的链接的颜色

/**
* Make a new cell with given key, value, and parent, and with
* {@code null} child links, and BLACK color.
*/

Entry(K key, V value, Entry<K,V> parent) {
this.key = key;
this.value = value;
this.parent = parent;
}

/**
* Returns the key.
*
* @return the key
*/

public K getKey() {
return key;
}

/**
* Returns the value associated with the key.
*
* @return the value associated with the key
*/

public V getValue() {
return value;
}

/**
* Replaces the value currently associated with the key with the given
* value.
*
* @return the value associated with the key before this method was
* called
*/

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;
//Key 和 Value都要 equals
return valEquals(key,e.getKey()) && valEquals(value,e.getValue());
}

//哈希值的计算,Key和Value的哈希值进行位异或
public int hashCode() {
int keyHash = (key==null ? 0 : key.hashCode());
int valueHash = (value==null ? 0 : value.hashCode());
return keyHash ^ valueHash;
}

public String toString() {
return key + "=" + value;
}
}

添加及更新操作

为了维持有序,添加及更新的代价较高,复杂度为 O(log(n)) 。插入节点后需要修复红黑树,使其恢复平衡状态,该操作在此不作介绍。

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
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); //Key 存在,更新value
} while (t != null);
}
else { //比较器为null,Key 必须实现 Comparable 接口
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); //Key 存在,更新value
} while (t != null);
}
//Key 不存在,新建节点,插入二叉树
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;
}

删除

从红黑树中删除一个节点比插入更为复杂,这里不作展开。

1
2
3
4
5
6
7
8
9
public V remove(Object key) {
Entry<K,V> p = getEntry(key); //先查找该节点
if (p == null)
return null;

V oldValue = p.value;
deleteEntry(p); //删除节点
return oldValue;
}
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
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

// Start fixup at replacement node, if it exists.
// 左子节点存在,则 replacement 为左子节点,否则为右子节点
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) //p 就是根节点
root = replacement;
else if (p == p.parent.left)//p 是父节点的左子节点
p.parent.left = replacement;
else//p 是父节点的右子节点
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;
}
}
}

查找

红黑树也是排序二叉树,按照排序二叉树的查找方法进行查找。复杂度为 O(log(n)) 。

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
public V get(Object key) {
Entry<K,V> p = getEntry(key);
return (p==null ? null : p.value);
}

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;
}

//使用比较器进行查找
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;
}

判断是否包含 key 或 value :

1
2
3
4
5
6
7
8
9
10
11
public boolean containsKey(Object key) {
return getEntry(key) != null;
}

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;
}

导航方法

NaviableMap 接口支持一系列的导航方法,有 firstEntry()、 lastEntry()、 lowerEntry()、 higherEntry()、 floorEntry()、 ceilingEntry()、 pollFirstEntry() 、 pollLastEntry() 等,它们的实现原理都是类似的,区别在于如何在排序的二叉树中查找到对应的节点。

以 lowerEntry() 和 floorEntry() 为例:

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
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
//小于给定的Key
public Map.Entry<K,V> lowerEntry(K key) {
return exportEntry(getLowerEntry(key));
}

final Entry<K,V> getLowerEntry(K key) {
Entry<K,V> p = root;
while (p != null) {
int cmp = compare(key, p.key);
//1. 如果节点 p 小于 key
if (cmp > 0) {
//1.1 节点 p 有右子树,则在右子树中搜索
if (p.right != null)
p = p.right;
//1.2 节点 p 没有右子树,找到目标
else
return p;
//2. 节点 p 大于等于 key
} else {
//2.1 节点 p 有左子树,则在左子树中继续搜索
if (p.left != null) {
p = p.left;
//2.2 节点 p 无左子树,找出 p 的前驱节点,并返回
//前驱节点要么不存在,要么就是小于 key 的最大节点
//因为从根节点一直遍历到 p,那么之前经过的所有节点都是大于等于 key 的
//且 p 没有左子树,即 p 是大于等于 key 的所有节点中最小的
//则 p 的前驱一定是查找的目标
} 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;
}
}
}
return null;
}

public K lowerKey(K key) {
return keyOrNull(getLowerEntry(key));
}

//小于等于
public Map.Entry<K,V> floorEntry(K key) {
return exportEntry(getFloorEntry(key));
}

//和 getLowerEntry 类似,相等时的处理不同
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;
} 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;
}

查找的过程可以和前驱节点的方法进行类比。 TreeMap 并没有直接暴露 getLowerEntry() 方法,而是使用 exportEntry(getLowerEntry(key)) 进行了一次包装。看似“多此一举”,实际上是为了防止对节点进行修改。SimpleImmutableEntry 类可以看作不可修改的 Key-Value 对,因为成员变量 key 和 value 都是 final 的。

即通过暴露出来的接口 firstEntry()、 lastEntry()、 lowerEntry()、 higherEntry()、 floorEntry()、 ceilingEntry() 是不可以修改获取的节点的,否则会抛出异常。

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
/**
* Return SimpleImmutableEntry for entry, or null if null
*/

static <K,V> Map.Entry<K,V> exportEntry(TreeMap.Entry<K,V> e) {
return (e == null) ? null :
new AbstractMap.SimpleImmutableEntry<>(e);
}

//AbstractMap.SimpleImmutableEntry
public static class SimpleImmutableEntry<K,V>
implements Entry<K,V>, java.io.Serializable
{

private static final long serialVersionUID = 7138329143949025153L;

private final K key;
private final V value;

public SimpleImmutableEntry(Entry<? extends K, ? extends V> entry) {
this.key = entry.getKey();
this.value = entry.getValue();
}

public V setValue(V value) {
throw new UnsupportedOperationException();
}
//....
//
}

pollFirstEntry() 、 pollLastEntry() 获取第一个和最后一个节点,并将它们从红黑树中删除。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
public Map.Entry<K,V> pollFirstEntry() {
Entry<K,V> p = getFirstEntry();
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();
Map.Entry<K,V> result = exportEntry(p);
if (p != null)
deleteEntry(p);
return result;
}

遍历

可以按照键的顺序遍历对 TreeSet 进行遍历,因为底层使用了红黑树来保证有序性,迭代器的实现就是按序访问排序二叉树中的节点。

先看一些内部抽象类 PrivateEntryIterator ,它是 TreeMap 中所有迭代器的基础:

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
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;
}

final Entry<K,V> prevEntry() {
Entry<K,V> e = next;
if (e == null)
throw new NoSuchElementException();
if (modCount != expectedModCount)
throw new ConcurrentModificationException();
next = predecessor(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;
}
}

因为红黑树自身就是有序的,迭代是只要从第一个节点不断获取后继节点即可。当然,逆序时则是从最后一个节点不断获取前驱节点。通过迭代器访问时基于 modCount 实现对并发修改的检查。

在排序二叉树中获取前驱和后继节点的方法如下:

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
//后继节点
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 {
//右子树不存在
//若父节点为null,则该节点是最大节点(根节点,且无右子树),无后继,返回null
//若当前节点是父节点的左子节点,直接返回父节点
//若当前节点是父节点的右子节点,则当前节点是以父节点为根的子树中最大的节点
Entry<K,V> p = t.parent; //父节点
Entry<K,V> ch = t;//当前节点
while (p != null && ch == p.right) {
//是右子节点,向上迭代,直到是左子节点
ch = p;
p = p.parent;
}
return p;
}
}

//前驱节点,同后继节点处理逻辑一致,左右颠倒
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;
}
}

其它方法

TreeMap 中还实现了一些其它的方法,如区间操作: headMap(), tailMap(), subMap() ; 获取逆序的 map: descendingMap() , descendingKeySet() 。只要了解了前面介绍的各种操作的原理,再来看这些方法的实现应该也不难理解。由于篇幅太长,这里就不再介绍了。

小结

TreeMap 是基于红黑树实现的一种 Key-Value 结构,最大的特点在于可以按照 Key 的顺序进行访问,要求 Key 实现 Comparable 接口或传入 Comparator 作为比较器。因为基于红黑树实现,TreeMap 内部在实现插入和删除操作时代价较高。

TreeMap 实现了 NavigableMap 接口,可以支持一系列导航方法,有 firstEntry()、 lastEntry()、 lowerEntry()、 higherEntry()、 floorEntry()、 ceilingEntry()、 pollFirstEntry() 、 pollLastEntry() ;还可以支持区间操作获取 map 的一部分,如 subMap(), headMap(), tailMap(K fromKey) 。除此以外, TreeMap 还支持通过 descendingMap() 获取和原来顺序相反的 map。

如果 TreeMap 没有使用自定义的 Comparator,则是不支持键为 null 的,因为调用 compareTo() 可能会发生异常;如果自定义的比较器可以接受 null 作为参数,那么是可以支持将 null 作为键的。

TreeMap 不是线程安全的,多线程情况下要手动进行同步或使用 SortedMap m = Collections.synchronizedSortedMap(new TreeMap(...));

-EOF-