08-TreeMap 核心源码解析（集合）

注：源码系列文章主要是对某付费专栏的总结记录。如有侵权，请联系删除。

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

从源码中，我们可以看到：

1. 新增节点时，就是利用了红黑树左小右大的特性，从根节点不断往下查找，直到找到节点是 null 为止，节点为 null 说明到达了叶子节点；
2. 查找的过程中，发现 key 值已经存在，直接覆盖值；
3. TreeMap 是禁止 key 是 null 值的。

类似的，TreeMap 查找也是类似的原理。

2.3 小结

TreeMap 相对来说比较简单，红黑树和 HashMap 比较类似，比较关键的是通过 compare 来比较 key 的大小，然后利用红黑树左小右大的特性，为每个 key 找到自己的位置，从而维护了 key 的大小排序顺序。

------------------------------------- END -------------------------------------