TreeMap源码摘要

发表于 | 分类于 | 阅读次数:
字数统计: 828 | 阅读时长 ≈ 4

本篇主要分析TreeMap源码。

TreeMap是用红黑树实现的，不妨按照之前《红黑树》的那一篇笔记的流程来分析，这里采用的是JDK1.8.0_91对应的源码。

初始化

注释写到：“Node in the Tree. Doubles as a means to pass key-value pairs back to user (see Map.Entry).”之前某篇我提到过JDK8在部分源码中已经使用Node替代Entry，这里仍然使用的是Entry：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
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;
    }

和以往不同的是，这里增加了parent字段，Entry是静态类。

Put

先看看源码中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
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;
}

这段代码一开始采用了compare(key, key)的比较，其实不难理解：如果root为null的话，构造root的时候(也就是new一个新的Entry)天然缺少类型检查，比如你可以不强迫新的Entry键值为类型K和V，而当root存在的时候，put()的元素都是会和父节点比较的，这样就隐式进行过类型(包括null)检查了。这里compare()是内部实现的一个final方法：

1
2
3
4
5
@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);
}

comparator不为null就用Comparator比较，为null就用Comparable比较。Comparator和Comparable的区别我在某篇中提到过，前者是外部比较器，后者可以视为内部比较器，通过重写外部比较器，可以方便更改比较规则。

modCount++的意涵在之前某源码摘要中也提到过，这里和那篇中的用法一样。fixAfterInsertion(e)用来插入后再平衡，基本原理见之前涉及红黑树的那篇。

应用

TreeMap的构造函数如下：

1
2
3
4
5
6
7
public TreeMap() {
    comparator = null;
}
public TreeMap(Comparator<? super K> comparator) {
    this.comparator = comparator;
}

可见，默认采用Comparable，Comparator实例可以作为实参更改排序规则。

Collections的一个静态方法reverseComparator()返回逆序比较器：

1
2
3
4
@SuppressWarnings("unchecked")
public static <T> Comparator<T> reverseOrder() {
    return (Comparator<T>) ReverseComparator.REVERSE_ORDER;
}

所以通过给构造器传入该方法，遍历的时候可以得到逆序结果。

再者，TreeMap的排序默认是高位优先(具体可见《算法(第4版)》)的。所以排序的内容如果有自己特定的逻辑规则，也应该自行实现Comparator。

小结

较之HashMap，TreeMap很稳，不然在Java8的HashMap源码中，就不会将某些特定情况转换成TreeMap来实现，基本操作的时间复杂度和红黑树一样。