Java 中的 Map 是一種鍵值對對映，又被稱為符號表或字典的資料結構，通常使用雜湊表來實現，但也可使用二叉查詢樹、紅黑樹實現。
- HashMap 基於雜湊表，但迭代時不是插入順序
- LinkedHashMap 擴充套件了 HashMap，維護了一個貫穿所有元素的雙向連結串列，保證按插入順序迭代
- TreeMap 基於紅黑樹，保證鍵的有序性，迭代時按鍵大小的排序順序
這裡就來分析下 TreeMap 的實現。基於紅黑樹，就意味著結點的增刪改查都能在 O(lgn) 時間複雜度內完成，如果按樹的中序遍歷就能得到一個按 鍵-key 大小排序的序列。
在看本文之前，建議看一下《紅黑樹這個資料結構，讓你又愛又恨？看了這篇，妥妥的征服它》對紅黑樹的分析，理解了紅黑樹，你會發現 TreeMap 如此簡單。
基本結構
TreeMap 的繼承結構如下，其中包含了一些關鍵欄位和方法：
其中，相關欄位的意義是：
- Comparator - 不為空，那麼就用它維持 key-鍵 的有序，否則使用 key-鍵 的自然順序
- size - 記錄樹中結點的個數
- modCount - 記錄樹結構變化次數，用於迭代器的快速失敗
另一個欄位是 Entry<K,V> root ，它表示根結點，初始為空，樹結點的結構定義如下：
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) { ... }
...
} TreeMap 是按照演算法導論（CLR）的描述實現的，但略有不同，它沒有使用隱形葉子結點 NIL，而是定義了一組訪問方法來正確處理 NULL 葉子節點 的問題，用於避免在主演算法中因檢查空葉子結點引起的混亂，方法如下：
- colorOf(Entry<K,V> p): 返回結點顏色，如果為空返回黑色
- parentOf(Entry<K,V> p): 返回父結點的引用，根結點則返回 null
- setColor(Entry<K,V> p, boolean c): 設定結點顏色
- leftOf(Entry<K,V> p): 返回左孩子結點
- rightOf(Entry<K,V> p): 返回右孩子結點
- rotateLeft(Entry<K,V> p): 將結點 P 左旋轉
- rotateRight(Entry<K,V> p): 將結點 P 右旋轉
- fixAfterInsertion(Entry<K,V> x): 插入結點後的回撥方法，重新平衡
- fixAfterDeletion(Entry<K,V> x): 刪除結點後的回撥方法，重新平衡
這些方法基本上都能見名知意，其中有點繞的就是樹旋轉的程式碼，程式碼實現如下：
插入
結點的插入可能會打破紅黑樹的平衡，需要做旋轉和顏色變換的調整。假設待插入結點為 N，P 是 N 的父結點，G 是 N 的祖父結點，U 是 N 的叔叔結點（即父結點的兄弟結點），那麼紅黑樹有以下幾種插入情況：
- N 是根結點，即紅黑樹的第一個結點
- N 的父結點（P）為黑色
- P 是紅色的（不是根結點），它的兄弟結點 U 也是紅色的
- P 為紅色，而 U 為黑色
4.1 P 左（右）孩子 N 右（左）孩子
4.2 P 左（右）孩子 N 左（右）孩子
以上情況的分析可檢視本文開頭的文章連結，現在來看下 TreeMap 的 put 方法的實現：
public V put(K key, V value) {
Entry<K,V> t = root;
// 情況 1 - 空樹，直接插入作為根結點
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) { // 使用 comparator 比較大小
do { // 根據 key 的大小找到插入位置
parent = t;
cmp = cpr.compare(key, t.key);
if (cmp < 0) t = t.left;
else if (cmp > 0) t = t.right;
else // 如果有相等的 key 直接設定 value 並返回
return t.setValue(value);
} while (t != null);
}
else {// 使用 key 的自然順序
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);
} // 新建一個結點插入
Entry<K,V>e = new Entry<>(key, value, parent);
if (cmp < 0) parent.left = e;
else parent.right = e;
fixAfterInsertion(e);// 可能會打破平衡，調整樹結構
size++;
modCount++;
return null;
}put 方法比較簡單，就是根據 key 的大小，遞迴的判斷插入左子樹還是右子樹，比較複雜操作在於插入後重新平衡的調整，核心程式碼如下：
刪除
結點的刪除也可能會打破紅黑樹的平衡，相比插入它的情況更復雜，假設待刪除結點為 M，如果有非葉子結點，稱為 C，那麼有兩種比較簡單的刪除情況：
- M 為紅色結點，那麼它必是葉子結點，直接刪除即可，因為如果它有一個黑色的非葉子結點，那麼就違反了性質5，通過 M 向左或向右的路徑黑色結點不等
- M 是黑色而 C 是紅色，只需要讓 C 替換到 M 的位置，並變成黑色即可，或者說交換 C 和 M 的值，並刪除 C（就是第一個簡單的情況）
這兩個情況，本質都是刪除了一個紅色結點，不影響整體平衡，比較複雜的是 M 和 C 都是黑色的情況，需要找一個結點填補這個黑色空缺。
結點 M刪除後它的位置上就變成了 NIL 隱形結點，為了方便描述，這個結點記為 N，P 表示 N 的父結點，S 表示 N 兄弟結點，S 如果存在左右孩子，分別使用 SL 和 SR 表示，那麼刪除就有以下幾種情況：
- N 是根結點 - 直接刪除即可
- P 黑 S 紅 - 交換 P 和 S 的顏色，然後對 P 左旋轉
- P 黑 S 黑 - 將 S 變成紅色，問題遞迴到父結點處理
- P 紅 S 黑 - 將 S 變成紅色，刪除成功
- P 顏色任意 S 黑 SL 紅 - 對 S 右旋轉，並交換 S 和 SL 的顏色，變成情況6
- P 顏色任意 S 黑，SR 紅 - 對 P 左旋轉，交換 P 和 S 的顏色，並將 SR 變成黑色
針對這些情況，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--;
// 如果 p 有兩個孩子結點，轉成刪除最多有一個孩子的結點的情況
// 這裡查詢的是 p 的後繼結點，也就是右子樹值最小的結點
if (p.left != null && p.right != null) {
Entry<K,V> s = successor(p); // 查詢後繼結點
// 複製後繼結點的 key 和 value 到 p
p.key = s.key;
p.value = s.value;
p = s; // 將 p 指向這個右子樹值最小的結點
} // p has 2 children
// 此時刪除的 p 要麼是葉子結點，要麼只有一個左或右孩子
Entry<K,V> replacement = (p.left != null ? p.left : p.right);
if (replacement != null) { // 有孩子結點
// 有一個左或右孩子，使用這個孩子結點替換它的父結點 p
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，也就是斷開所有的連結
p.left = p.right = p.parent = null;
// Fix replacement. 如果刪除的是黑色結點
if (p.color == BLACK)
fixAfterDeletion(replacement); // 平衡調整
} else if (p.parent == null) { // return if we are the only node.
root = null;// 情況1，刪除後變成空樹
} else {//No children. Use self as phantom replacement and unlink.
// 刪除的是葉子結點，那麼刪除 p 就是用它的隱形 NIL 葉子結點替換
// 它，這裡將它自己看做隱形的葉子結點
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;
}
}
}deleteEntry 的邏輯就和二叉查詢樹一樣，主要就是把刪除任一結點的問題就簡化成：刪除一個最多隻有一個孩子的結點的情況，並且所有的刪除操作都在葉子結點完成。如果刪除的是黑色結點，那麼就視情況調整樹重新達到平衡，具體程式碼如下:
查詢
就像二分查詢那樣，TreeMap 也能在 ~lgN 次比較內結束查詢，並且針對 鍵-key 提供了豐富的查詢 API，
- get(Object key) - 返回等於給定鍵的結點
- floorEntry(K key) - 返回小於或等於給定鍵的結點中鍵最大的結點
- ceilingEntry(K key) - 返回大於或等於給定鍵的結點中鍵最小的結點
- higherEntry(K key) - 返回嚴格大於給定鍵的結點中鍵最小的結點
- lowerEntry(K key) - 返回嚴格小於給定鍵的結點中鍵最大的結點
上面這些方法比較簡單，可自行檢視原始碼。另外，還有兩個比較特殊的方法，它們用來查詢指定結點在樹中序遍歷序列中的前驅和後繼結點，在中序遍歷序列中：
- 前驅結點也就是左子樹值最大的結點
- 後繼結點也就是右子樹值最小的結點
遍歷
遍歷也是一個高頻操作，在 Java 集合框架體系中，基本都是採用迭代器 Iterator 來實現，TreeMap 也是如此，它提供了對鍵和對值的迭代器。
TreeMap 迭代器最終的邏輯實現是在 PrivateEntryIterator 類中，預設按鍵的正序輸出，它也提供了一個逆序輸出的迭代器 DescendingKeyIterator。
具體程式碼不在貼出，比較簡單，值得注意的就是上一節介紹的查詢前驅和後繼結點的兩個方法，遍歷常用 API 有：
- entrySet() - 返回一個遍歷所有結點的 Set 集合
- keySet() - 返回一個遍歷所有鍵的 Set 集合
- values() - 返回一個遍歷所有值的 Set 集合
小結
分析 TreeMap 的原始碼之前，一定要去分析紅黑樹的原理，然後在看它的原始碼，相信理論與實踐相結合，掌握紅黑樹不在話下，TreeMap 也會用得遊刃有餘。