TreeMap源码学习

发布于 — 2019 年 11 月 12 日
#Java

之前看过了HashMap,LinkedHashMap的源码,这次来看一下TreeMap的源码。

从这个名字就能看出,TreeMap底层使用的是树来进行存储的。

变量

1
2
3
4
5
6
7
8
9
//比较器,用于左右子树的判断。
//正常情况下,左子树为 1,父节点为 2,右子树为 3。如果比较器设置 3<1<2。则左子树为3,父节点为 1,右子树为 2。
private final Comparator<? super K> comparator;
//根节点
private transient Entry<K,V> root;
//容量
private transient int size = 0;
//修改的次数,在迭代和序列化时用到
private transient int modCount = 0;

看一下 root 节点的数据结构:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
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;

	...
}

由于有一个color=BLACK属性,所以底层数据结构应该是红黑树

构造器

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
// 无参构造器
public TreeMap() {
    comparator = null;
}
//传入比较器
public TreeMap(Comparator<? super K> comparator) {
    this.comparator = comparator;
}
//传入 map
public TreeMap(Map<? extends K, ? extends V> m) {
    comparator = null;
    putAll(m);
}
//传入一个排序的 map
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) {
    }
}

get

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

在上面代码中的第八行进行了判断comparator也就是自定义的比较器是否为空,这两种情况下查找的比较器和 key 有所不同。然后再进行二分查找。

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;
    if (t == null) {
        //当 map 为空时
        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 {
            //对树进行插入,如果key 重复则更新 value
            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;
}

首先进行了 map 为空时的判断。然后对比较器为空和非空时的逻辑。最后调用了对红黑树重排的方法。