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