08-TreeMap 核心源码解析(集合)

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

1 知识储备

在了解 TreeMap 之前,我们来看看日常工作中排序的两种方式,作为我们学习的基础储备,两种方式的代码如下:

@Data
class Entry implements Comparable {
    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 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 list2 = new ArrayList<>();
    for (int i = 5; i > 0; i--) {
        list2.add(new Entry(i));
    }
    Collections.sort(list2, new Comparator() {
        @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 comparator;

// 红黑树的根节点
private transient Entry 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 implements Map.Entry {
    K key;
    V value;
    Entry left;
    Entry right;
    Entry parent;
    boolean color = BLACK;

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

2.2 新增节点

TreeMap 新增节点源码:

public V put(K key, V value) {
    Entry 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 parent;
    // split comparator and comparable paths
    Comparator 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 k = (Comparable) 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 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 -------------------------------------

你可能感兴趣的:(08-TreeMap 核心源码解析(集合))