TreeMap 是 Map 集合的有序實現，其底層是基於紅黑樹的實現，能夠早 log(n) 時間內完成 get、put 和 remove 操作。
public class TreeMap<K,V>
extends AbstractMap<K,V>
implements NavigableMap<K,V>, Cloneable, java.io.Serializable
TreeMap 繼承自 AbstractMap，還實現了 NavigableMap接口。NavigableMap 接口繼承了SortedMap接口，SortedMap 最終繼承自Map接口。整體來說 TreeMap 的繼承體系如下圖所示。
原理
我們將從類成員變量、構造方法、核心方法幾個方面來解析 TreeMap 的實現。
類成員變量
// 比較器。根據這個比較器決定TreeMap的排序。
// 如果爲空，表示按照key做自然排序（最小的在根節點）。
private final Comparator<? super K> comparator;
// 根節點
private transient Entry<K,V> root;
// 大小
private transient int size = 0;
// Node節點聲明
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;
}
我們可以看到 TreeMap 有一個 Entry 類型的 root 節點，而 Entry 則是 TreeMap 的內部類。從 TreeMap.Entry 的屬性我們可以知道其實一個紅黑樹節點的實現。
構造方法
TeeMap 一共有四個構造方法：
public TreeMap() {
comparator = null;
}
public TreeMap(Comparator<? super K> comparator) {
this.comparator = comparator;
}
public TreeMap(Map<? extends K, ? extends V> m) {
comparator = null;
putAll(m);
}
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) {
}
}
核心方法
我們將從查找、插入、刪除、遍歷四個方法研究 TreeMap 的實現。
查找
TreeMap基於紅黑樹實現，而紅黑樹是一種自平衡二叉查找樹，所以 TreeMap 的查找操作流程和二叉查找樹一致。二叉樹的查找流程是這樣的，先將目標值和根節點的值進行比較，如果目標值小於根節點的值，則再和根節點的左孩子進行比較。如果目標值大於根節點的值，則繼續和根節點的右孩子比較。在查找過程中，如果目標值和二叉樹中的某個節點值相等，則返回 true，否則返回 false。
TreeMap 查找和此類似，只不過在 TreeMap 中，節點（Entry）存儲的是鍵值對<k,v>。在查找過程中，比較的是鍵的大小，返回的是值，如果沒找到，則返回null。TreeMap 中的查找方法是get，具體實現在getEntry方法中，相關源碼如下：
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;
}
插入
TreeMap 的插入其實就是紅黑樹的插入，因此搞懂了紅黑樹插入的各個情況，看懂 TreeMap 的插入源碼就不在話下了。
public V put(K key, V value) {
Entry<K,V> t = root;
// 1. 如果根節點爲 null，將新節點設爲根節點
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;
// 2.爲 key 在紅黑樹找到合適的位置
Comparator<? super K> cpr = comparator;
if (cpr != null) {
do {
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);
}
// 3.將新節點鏈入紅黑樹中
Entry<K,V> e = new Entry<>(key, value, parent);
if (cmp < 0)
parent.left = e;
else
parent.right = e;
// 4.插入新節點可能會破壞紅黑樹性質，這裏修正一下
fixAfterInsertion(e);
size++;
modCount++;
return null;
}
刪除
TreeMap 的刪除其實就是紅黑樹的刪除，因此搞懂了紅黑樹刪除的各個情況，看懂 TreeMap 的刪除源碼就不在話下了。
public V remove(Object key) {
Entry<K,V> p = getEntry(key);
if (p == null)
return null;
V oldValue = p.value;
deleteEntry(p);
return oldValue;
}
private void deleteEntry(Entry<K,V> p) {
modCount++;
size--;
/*
* 1. 如果 p 有兩個孩子節點，則找到後繼節點，
* 並把後繼節點的值複製到節點 P 中，並讓 p 指向其後繼節點
*/
if (p.left != null && p.right != null) {
Entry<K,V> s = successor(p);
p.key = s.key;
p.value = s.value;
p = s;
} // p has 2 children
// Start fixup at replacement node, if it exists.
Entry<K,V> replacement = (p.left != null ? p.left : p.right);
if (replacement != null) {
/*
* 2. 將 replacement parent 引用指向新的父節點，
* 同時讓新的父節點指向 replacement。
*/
replacement.parent = p.parent;
if (p.parent == null)
root = replacement;
else if (p == p.parent.left)
p.parent.left = replacement;
else
p.parent.right = replacement;
// Null out links so they are OK to use by fixAfterDeletion.
p.left = p.right = p.parent = null;
// 3. 如果刪除的節點 p 是黑色節點，則需要進行調整
if (p.color == BLACK)
fixAfterDeletion(replacement);
} else if (p.parent == null) { // 刪除的是根節點，且樹中當前只有一個節點
root = null;
} else { // 刪除的節點沒有孩子節點
// p 是黑色，則需要進行調整
if (p.color == BLACK)
fixAfterDeletion(p);
// 將 P 從樹中移除
if (p.parent != null) {
if (p == p.parent.left)
p.parent.left = null;
else if (p == p.parent.right)
p.parent.right = null;
p.parent = null;
}
}
}
遍歷
TreeMap 有一個特性，即可以保證鍵的有序性，默認是正序。所以在遍歷過程中，大家會發現 TreeMap 會從小到大輸出鍵的值。那麼，接下來就來分析一下keySet方法，以及在遍歷 keySet 方法產生的集合時，TreeMap 是如何保證鍵的有序性的。相關代碼如下：
public Set<K> keySet() {
return navigableKeySet();
}
public NavigableSet<K> navigableKeySet() {
KeySet<K> nks = navigableKeySet;
return (nks != null) ? nks : (navigableKeySet = new KeySet<>(this));
}
static final class KeySet<E> extends AbstractSet<E> implements NavigableSet<E> {
private final NavigableMap<E, ?> m;
KeySet(NavigableMap<E,?> map) { m = map; }
public Iterator<E> iterator() {
if (m instanceof TreeMap)
return ((TreeMap<E,?>)m).keyIterator();
else
return ((TreeMap.NavigableSubMap<E,?>)m).keyIterator();
}
// 省略非關鍵代碼
}
Iterator<K> keyIterator() {
return new KeyIterator(getFirstEntry());
}
final class KeyIterator extends PrivateEntryIterator<K> {
KeyIterator(Entry<K,V> first) {
super(first);
}
public K next() {
return nextEntry().key;
}
}
abstract class PrivateEntryIterator<T> implements Iterator<T> {
Entry<K,V> next;
Entry<K,V> lastReturned;
int expectedModCount;
PrivateEntryIterator(Entry<K,V> first) {
expectedModCount = modCount;
lastReturned = null;
next = first;
}
public final boolean hasNext() {
return next != null;
}
final Entry<K,V> nextEntry() {
Entry<K,V> e = next;
if (e == null)
throw new NoSuchElementException();
if (modCount != expectedModCount)
throw new ConcurrentModificationException();
// 尋找節點 e 的後繼節點
next = successor(e);
lastReturned = e;
return e;
}
// 其他方法省略
}
上面的代碼比較多，keySet 涉及的代碼還是比較多的，大家可以從上往下看。從上面源碼可以看出 keySet 方法返回的是 KeySet 類的對象。這個類實現了Iterable接口，可以返回一個迭代器。該迭代器的具體實現是 KeyIterator，而 KeyIterator 類的核心邏輯是在PrivateEntryIterator中實現的。上面的代碼雖多，但核心代碼還是 KeySet 類和 PrivateEntryIterator 類的 nextEntry方法。KeySet 類就是一個集合，這裏不分析了。而 nextEntry 方法比較重要，下面簡單分析一下。
在初始化 KeyIterator 時，默認情況下會將 TreeMap 中包含最小鍵或最大值（取決於傳入的比較器）的 Entry 傳給 PrivateEntryIterator。當調用 nextEntry 方法時，通過調用 successor 方法找到當前 entry 的後繼，並讓 next 指向後繼，最後返回當前的 entry。通過這種方式即可實現按正序返回鍵值的的邏輯。
總結
TreeMap 是 Map 集合的經典紅黑樹實現，所以弄懂了紅黑樹的查詢、插入、刪除，TreeMap 的源碼自然不再話下。但要注意的是，TreeMap 的遍歷並不是從根節點開始遍歷，而是根據 key 的大小從小到大輸出，或者從大到小輸出。到底是升序還是降序，取決於傳入的 Comparator，默認是升序，即從小到大輸出。
TreeMap 是哈希的紅黑樹經典實現，是Map的哈希有序實現。