TreeMap概览

相对于HashMap来说，TreeMap 是较简单的。

类介绍（注释）

1. TreeMap是红黑树的实现，它可以通过自然排序、构造时传入的Comparator对key进行排序。
2. TreeMap（和其他的可排序的Map）中的元素，在没显式传入构造器的时候，一定要使该类 实现Comparable、equals方法。
3. containsKey、get、 put、remove 方法，时间复杂度是O(log(n))
4. TreeMap是非线程安全的，并发出错时，会快速失败，抛出ConcurrentModificationException
   

TreeMap的一些概念

TreeMap 底层的数据结构就是红黑树，和 HashMap 的红黑树结构一样。
不同的是，TreeMap 利用了红黑树左节点小，右节点大的性质，根据 key 进行排序，使每个元素能够插入到红黑树大小适当的位置，维护了 key 的大小关系，适用于 key 需要排序的场景。

//比较器，如果外部有传进来 Comparator 比较器，首先用外部的
//如果外部比较器为空，则使用 key 自己实现的 Comparable#compareTo 方法
private final Comparator<? super K> comparator;

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

//红黑树的 已有元素大小
private transient int size = 0;

//树结构变化的版本号
private transient int modCount = 0;

//红黑树的节点
static final class Entry<K,V> implements Map.Entry<K,V> {}

TreeMap的使用Demo

public static void main(String[] args) {
    // 1. 类实现compareTo
    TreeMap<Person, Object> map = new TreeMap<>();
	// 2. 显式传入比较器
    TreeMap<Person, Object> map2 = new TreeMap<>(new Comparator<Person>() {
        @Override
        public int compare(Person o1, Person o2) {
            return o1.age - o2.age;
        }
    });

    for (int i = 0; i < 5; i++) {
        Person person = new Person(i);
        map.put(person, i);
        map2.put(person, i);
    }
    map.forEach((k, v) -> System.out.println(k));
    System.out.println("============");
    map2.forEach((k, v) -> System.out.println(k));

}

@Data
static class Person implements Comparable<Person> {
    private int age;
    private String name;

    public Person(int age) {
        this.age = age;
    }

    @Override
    public int compareTo(Person o) {

        return o.age - this.age;
    }

    @Override
    public boolean equals(Object o) {
        if (this == o) return true;
        if (o == null || getClass() != o.getClass()) return false;
        Person person = (Person) o;
        return age == person.age &&
            Objects.equals(name, person.name);
    }

    @Override
    public int hashCode() {
        return Objects.hash(age, name);
    }
}

运行结果

TreeMapTest.Person(age=4, name=null)
TreeMapTest.Person(age=3, name=null)
TreeMapTest.Person(age=2, name=null)
TreeMapTest.Person(age=1, name=null)
TreeMapTest.Person(age=0, name=null)
============
TreeMapTest.Person(age=0, name=null)
TreeMapTest.Person(age=1, name=null)
TreeMapTest.Person(age=2, name=null)
TreeMapTest.Person(age=3, name=null)
TreeMapTest.Person(age=4, name=null)

常用方法源码

put

put 方法大致分为三步骤。

1. 红黑树为空时（根节点为null），直接将put的元素，作为根节点元素。

// 备份根节点
Entry<K,V> t = root;
// 根节点不存在
if (t == null) {
    // 通过compare方法，比较key,(优先使用显式传入的比较器)；
    // 该方法同时限制了key不为null
    compare(key, key); // type (and possibly null) check
	// 创建根节点（入参 parent 为null）
    root = new Entry<>(key, value, null);
    
    // 更新树大小、结构版本号
    size = 1;
    modCount++;
    return null;
}

2. 根据红黑树左小右大的特性，进行判断，找到应该新增节点的父节点：（省略了561~575，因为逻辑相同，只不过用的比较器不同）

int cmp;
// 父节点
Entry<K,V> parent;

Comparator<? super K> cpr = comparator;
// 优先使用显式传入的比较器
if (cpr != null) {
    do {
        parent = t;
        // 比较
        cmp = cpr.compare(key, t.key);
        // 即key < t.key，因为左小右大，为了找到合适位置，则找更小的比较
        if (cmp < 0)
            t = t.left;
        
        // 即key > t.key，因为左小右大，为了找到合适位置，则找更大的比较
        else if (cmp > 0)
            t = t.right;
        
        // 即key = t.key，直接覆盖
        else
            return t.setValue(value);
        
        // 终止条件t==null，即t为叶子节点时
    } while (t != null);
}

3. 新节点成为父节点的左叶子或右叶子。以上循环中，我们知道parent永远是t的父节点。

// 创建新节点，其父节点是最后一次循环得到的 叶子节点
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;
// 修复红黑树（着色旋转，达到平衡）
fixAfterInsertion(e);
// 更新树大小、结构版本号
size++;
modCount++;