一、簡介
TreeMap最早出現在JDK 1.2中，是 Java 集合框架中比較重要一個的實現。TreeMap 底層基於紅黑樹實現，可保證在log(n)時間複雜度內完成 containsKey、get、put 和 remove 操作，效率很高。另一方面，由於 TreeMap 基於紅黑樹實現，這為 TreeMap 保持鍵的有序性打下了基礎。總的來說，TreeMap 的核心是紅黑樹，其很多方法也是對紅黑樹增刪查基礎操作的一個包裝。所以只要弄懂了紅黑樹，TreeMap 就沒什麼祕密了。
“
二、概覽
TreeMap繼承自AbstractMap，並實現了 NavigableMap介面。NavigableMap 介面繼承了SortedMap介面，SortedMap 最終繼承自Map介面，同時 AbstractMap 類也實現了 Map 介面。以上就是 TreeMap 的繼承體系，描述起來有點亂，不如看圖了：
上圖就是 TreeMap 的繼承體系圖，比較直觀。這裡來簡單說一下繼承體系中不常見的介面NavigableMap和SortedMap，這兩個介面見名知意。先說 NavigableMap 介面，NavigableMap 介面宣告瞭一些列具有導航功能的方法，比如：
/**
- 返回紅黑樹中最小鍵所對應的 Entry
*/
Map.EntryfirstEntry();
/**
- 返回最大的鍵 maxKey，且 maxKey 僅小於引數 key
*/
KlowerKey(K key);
/**
- 返回最小的鍵 minKey，且 minKey 僅大於引數 key
*/
KhigherKey(K key);
// 其他略
通過這些導航方法，我們可以快速定位到目標的 key 或 Entry。至於 SortedMap 介面，這個介面提供了一些基於有序鍵的操作，比如
/**
- 返回包含鍵值在 [minKey, toKey) 範圍內的 Map
*/
SortedMapheadMap(K toKey);();
/**
- 返回包含鍵值在 [fromKey, toKey) 範圍內的 Map
*/
SortedMapsubMap(K fromKey, K toKey);
// 其他略
以上就是兩個介面的介紹，很簡單。至於 AbstractMap 和 Map 這裡就不說了，大家有興趣自己去看看 Javadoc 吧。關於 TreeMap 的繼承體系就這裡就說到這，接下來我們進入細節部分分析。
“
三、原始碼分析
JDK 1.8中的TreeMap原始碼有兩千多行，還是比較多的。本文並不打算逐句分析所有的原始碼，而是挑選幾個常用的方法進行分析。這些方法實現的功能分別是查詢、遍歷、插入、刪除等，其他的方法小夥伴們有興趣可以自己分析。TreeMap實現的核心部分是關於紅黑樹的實現，其絕大部分的方法基本都是對底層紅黑樹增、刪、查操作的一個封裝。如簡介一節所說，只要弄懂了紅黑樹原理，TreeMap 就沒什麼祕密了。
TreeMap基於紅黑樹實現，而紅黑樹是一種自平衡二叉查詢樹，所以 TreeMap 的查詢操作流程和二叉查詢樹一致。二叉樹的查詢流程是這樣的，先將目標值和根節點的值進行比較，如果目標值小於根節點的值，則再和根節點的左孩子進行比較。如果目標值大於根節點的值，則繼續和根節點的右孩子比較。在查詢過程中，如果目標值和二叉樹中的某個節點值相等，則返回 true，否則返回 false。TreeMap 查詢和此類似，只不過在 TreeMap 中，節點（Entry）儲存的是鍵值對。在查詢過程中，比較的是鍵的大小，返回的是值，如果沒找到，則返回null。TreeMap 中的查詢方法是get，具體實現在getEntry方法中，相關原始碼如下：
publicVget(Object key) {
Entry p = getEntry(key);
return(p==null?null: p.value);
}
finalEntry getEntry(Object key) {
// Offload comparator-based version for sake of performance
if(comparator !=null)
returngetEntryUsingComparator(key);
if(key ==null)
thrownew NullPointerException();
@SuppressWarnings("unchecked")
Comparable k = (Comparable) key;
Entry p = root;
// 查詢操作的核心邏輯就在這個 while 迴圈裡
while(p !=null) {
int cmp = k.compareTo(p.key);
if(cmp <0)
p = p.left;
elseif(cmp >0)
p = p.right;
else
returnp;
}
returnnull;
}
查詢操作的核心邏輯就是getEntry方法中的while迴圈，大家對照上面的說的流程，自己看一下吧，比較簡單，就不多說了。
3.2 遍歷
遍歷操作也是大家使用頻率較高的一個操作，對於TreeMap，使用方式一般如下：
for(Object key :map.keySet()) {
// do something
}
或
for(Map.Entry entry :map.entrySet()) {
// do something
}
從上面程式碼片段中可以看出，大家一般都是對 TreeMap 的 key 集合或 Entry 集合進行遍歷。上面程式碼片段中用 foreach 遍歷keySet 方法產生的集合，在編譯時會轉換成用迭代器遍歷，等價於：
Setkeys = map.keySet();
Iterator ite = keys.iterator();
while(ite.hasNext()) {
Objectkey = ite.next();
// do something
}
另一方面，TreeMap 有一個特性，即可以保證鍵的有序性，預設是正序。所以在遍歷過程中，大家會發現 TreeMap 會從小到大輸出鍵的值。那麼，接下來就來分析一下keySet方法，以及在遍歷 keySet 方法產生的集合時，TreeMap 是如何保證鍵的有序性的。相關程式碼如下：
publicSetkeySet(){
returnnavigableKeySet();
}
publicNavigableSetnavigableKeySet(){
KeySet nks = navigableKeySet;
return(nks !=null) ? nks : (navigableKeySet =newKeySet<>(this));
}
staticfinalclassKeySetextendsAbstractSetimplementsNavigableSet{
privatefinalNavigableMap m;
KeySet(NavigableMap map) { m = map; }
publicIteratoriterator(){
if(minstanceofTreeMap)
return((TreeMap)m).keyIterator();
else
return((TreeMap.NavigableSubMap)m).keyIterator();
}
// 省略非關鍵程式碼
}
IteratorkeyIterator(){
returnnewKeyIterator(getFirstEntry());
}
finalclassKeyIteratorextendsPrivateEntryIterator{
KeyIterator(Entry first) {
super(first);
}
publicKnext(){
returnnextEntry().key;
}
}
abstractclassPrivateEntryIteratorimplementsIterator{
Entry next;
Entry lastReturned;
intexpectedModCount;
PrivateEntryIterator(Entry first) {
expectedModCount = modCount;
lastReturned =null;
next = first;
}
publicfinalbooleanhasNext(){
returnnext !=null;
}
finalEntrynextEntry(){
Entry e = next;
if(e ==null)
thrownewNoSuchElementException();
if(modCount != expectedModCount)
thrownewConcurrentModificationException();
// 尋找節點 e 的後繼節點
next = successor(e);
lastReturned = e;
returne;
}
// 其他方法省略
}
上面的程式碼比較多，keySet 涉及的程式碼還是比較多的，大家可以從上往下看。從上面原始碼可以看出 keySet 方法返回的是KeySet類的物件。這個類實現了Iterable介面，可以返回一個迭代器。該迭代器的具體實現是KeyIterator，而 KeyIterator 類的核心邏輯是在PrivateEntryIterator中實現的。上面的程式碼雖多，但核心程式碼還是 KeySet 類和 PrivateEntryIterator 類的 nextEntry方法。KeySet 類就是一個集合，這裡不分析了。而 nextEntry 方法比較重要，下面簡單分析一下。
在初始化 KeyIterator 時，會將 TreeMap 中包含最小鍵的 Entry 傳給 PrivateEntryIterator。當呼叫 nextEntry 方法時，通過呼叫 successor 方法找到當前 entry 的後繼，並讓 next 指向後繼，最後返回當前的 entry。通過這種方式即可實現按正序返回鍵值的的邏輯。
好了，TreeMap 的遍歷操作就講到這。遍歷操作本身不難，但講的有點多，略顯囉嗦，大家見怪。
3.3 插入
相對於前兩個操作，插入操作明顯要複雜一些。當往 TreeMap 中放入新的鍵值對後，可能會破壞紅黑樹的性質。這裡為了描述方便，把 Entry 稱為節點。並把新插入的節點稱為N，N 的父節點為P。P 的父節點為G，且 P 是 G 的左孩子。P 的兄弟節點為U。在往紅黑樹中插入新的節點 N 後（新節點為紅色），會產生下面5種情況：
1、N 是根節點；
2、N 的父節點是黑色；
3、N 的父節點是紅色，叔叔節點也是紅色；
4、N 的父節點是紅色，叔叔節點是黑色，且 N 是 P 的右孩子；
5、N 的父節點是紅色，叔叔節點是黑色，且 N 是 P 的左孩子。
上面5中情況中，情況2不會破壞紅黑樹性質，所以無需處理。情況1 會破壞紅黑樹性質2（根是黑色），情況3、4、和5會破壞紅黑樹性質4（每個紅色節點必須有兩個黑色的子節點）。這個時候就需要進行調整，以使紅黑樹重新恢復平衡。接下來分析一下插入操作相關原始碼：
publicVput(K key, Vvalue){
Entry t = root;
// 1.如果根節點為 null，將新節點設為根節點
if(t ==null) {
compare(key, key);
root =newEntry<>(key,value,null);
size =1;
modCount++;
returnnull;
}
intcmp;
Entry parent;
// split comparator and comparable paths
Comparator cpr = comparator;
if(cpr !=null) {
// 2.為 key 在紅黑樹找到合適的位置
do{
parent = t;
cmp = cpr.compare(key, t.key);
if(cmp <0)
t = t.left;
elseif(cmp >0)
t = t.right;
else
returnt.setValue(value);
}while(t !=null);
}else{
// 與上面程式碼邏輯類似，省略
}
Entry e =newEntry<>(key,value, parent);
// 3.將新節點鏈入紅黑樹中
if(cmp <0)
parent.left = e;
else
parent.right = e;
// 4.插入新節點可能會破壞紅黑樹性質，這裡修正一下
fixAfterInsertion(e);
size++;
modCount++;
returnnull;
}
put 方法程式碼如上，邏輯和二叉查詢樹插入節點邏輯一致。重要的步驟我已經寫了註釋，並不難理解。插入邏輯的複雜之處在於插入後的修復操作，對應的方法fixAfterInsertion，該方法的原始碼和說明如下：
到這裡，插入操作就講完了。接下來，來說說 TreeMap 中最複雜的部分，也就是刪除操作了。
3.4 刪除
刪除操作是紅黑樹最複雜的部分，原因是該操作可能會破壞紅黑樹性質5（從任一節點到其每個葉子的所有簡單路徑都包含相同數目的黑色節點），修復性質5要比修復其他性質（性質2和4需修復，性質1和3不用修復）複雜的多。當刪除操作導致性質5被破壞時，會出現8種情況。為了方便表述，這裡還是先做一些假設。我們把最終被刪除的節點稱為 X，X 的替換節點稱為 N。N 的父節點為P，且 N 是 P 的左孩子。N 的兄弟節點為S，S 的左孩子為 SL，右孩子為 SR。這裡特地強調 X 是 最終被刪除 的節點，是原因二叉查詢樹會把要刪除有兩個孩子的節點的情況轉化為刪除只有一個孩子的節點的情況，該節點是欲被刪除節點的前驅和後繼。
接下來，簡單列舉一下刪除節點時可能會出現的情況，先列舉較為簡單的情況：
1、最終被刪除的節點 X 是紅色節點；
2、X 是黑色節點，但該節點的孩子節點是紅色。
比較複雜的情況：
1、替換節點 N 是新的根；
2、N 為黑色，N 的兄弟節點 S 為紅色，其他節點為黑色；
3、N 為黑色，N 的父節點 P，兄弟節點 S 和 S 的孩子節點均為黑色；
4、N 為黑色，P 是紅色，S 和 S 孩子均為黑色；
5、N 為黑色，P 可紅可黑，S 為黑色，S 的左孩子 SL 為紅色，右孩子 SR 為黑色；
6、N 為黑色，P 可紅可黑，S 為黑色，SR 為紅色，SL 可紅可黑。
上面列舉的8種情況中，前兩種處理起來比較簡單，後6種情況中情況26較為複雜。接下來我將會對情況26展開分析，刪除相關的原始碼如下：
publicV remove(Object key) {
Entry p = getEntry(key);
if(p ==null)
returnnull;
V oldValue = p.value;
deleteEntry(p);
returnoldValue;
}
privatevoid deleteEntry(Entry p) {
modCount++;
size--;
/*
-
- 如果 p 有兩個孩子節點，則找到後繼節點，
-
並把後繼節點的值複製到節點 P 中，並讓 p 指向其後繼節點
*/
if(p.left !=null&& p.right !=null) {
Entry 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 replacement = (p.left !=null? p.left : p.right);
if(replacement !=null) {
/*
-
- 將 replacement parent 引用指向新的父節點，
-
同時讓新的父節點指向 replacement。
*/
replacement.parent= p.parent;
if(p.parent==null)
root = replacement;
elseif(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);
}elseif(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;
elseif(p == p.parent.right)
p.parent.right =null;
p.parent=null;
}
}
}
從原始碼中可以看出，remove方法只是一個簡單的保證，核心實現在deleteEntry方法中。deleteEntry 主要做了這麼幾件事：
1、如果待刪除節點 P 有兩個孩子，則先找到 P 的後繼 S，然後將 S 中的值拷貝到 P 中，並讓 P 指向 S；
2、如果最終被刪除節點 P（P 現在指向最終被刪除節點）的孩子不為空，則用其孩子節點替換掉；
3、如果最終被刪除的節點是黑色的話，呼叫 fixAfterDeletion 方法進行修復。
上面說了 replacement 不為空時，deleteEntry 的執行邏輯。上面說的略微囉嗦，如果簡單說的話，7個字即可總結：找後繼 -> 替換 -> 修復。這三步中，最複雜的是修復操作。修復操作要重新使紅黑樹恢復平衡，修復操作的原始碼分析如下：
fixAfterDeletion 方法分析如下：
上面對 fixAfterDeletion 部分程式碼邏輯就行了分析，通過配圖的形式解析了每段程式碼邏輯所處理的情況。通過圖解，應該還是比較好理解的。好了，TreeMap 原始碼先分析到這裡。
“
四、總結
本文則從實踐層面是分析了插入和刪除操作在具體的實現中時怎樣做的。另外，本文選擇了從集合框架常用方法這一角度進行分析，詳細分析了查詢、遍歷、插入和刪除等方法。總體來說，分析的還是比較詳細的。當然限於本人的水平，文中可能會存在一些錯誤的論述。如果大家發現了，歡迎指出來。如果這些錯誤的論述對你造成了困擾，我這裡先說聲抱歉。如果你也在學習 TreeMap 原始碼，希望這篇文章能夠幫到你。
最後感謝大家花時間的閱讀我的文章，順祝大家寫程式碼無BUG，下篇文章見。
如果您覺得不錯，請別忘了轉發、分享、點贊讓更多的人去學習，順便給大家推薦一個架構交流群：617434785，裡面會分享一些資深架構師錄製的視訊錄影：有Spring，MyBatis，Netty原始碼分析，高併發、高效能、分散式、微服務架構的原理，JVM效能優化這些成為架構師必備的知識體系。還能領取免費的學習資源。相信對於已經工作和遇到技術瓶頸的碼友，在這個群裡會有你需要的內容。