一、簡介
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.Entry<K,V> firstEntry(); /** * 傳回最大的鍵 maxKey，且 maxKey 僅小于參數 key */ K lowerKey(K key); /** * 傳回最小的鍵 minKey，且 minKey 僅大于參數 key */ K higherKey(K key); // 其他略 通過這些導航方法，我們可以快速定位到目标的 key 或 Entry。至于 SortedMap 接口，這個接口提供了一些基于有序鍵的操作，比如
/** * 傳回包含鍵值在 [minKey, toKey) 範圍内的 Map */ SortedMap<K,V> headMap(K toKey);(); /** * 傳回包含鍵值在 [fromKey, toKey) 範圍内的 Map */ SortedMap<K,V> subMap(K fromKey, K toKey); // 其他略 以上就是兩個接口的介紹，很簡單。至于 AbstractMap 和 Map 這裡就不說了，大家有興趣自己去看看 Javadoc 吧。關于 TreeMap 的繼承體系就這裡就說到這，接下來我們進入細節部分分析。
三、源碼分析
JDK 1.8
中的
TreeMap
源碼有兩千多行，還是比較多的。本文并不打算逐句分析所有的源碼，而是挑選幾個常用的方法進行分析。這些方法實作的功能分别是查找、周遊、插入、删除等，其他的方法小夥伴們有興趣可以自己分析。
TreeMap
實作的核心部分是關于
紅黑樹
的實作，其絕大部分的方法基本都是對底層紅黑樹增、删、查操作的一個封裝。如簡介一節所說，隻要弄懂了紅黑樹原理，TreeMap 就沒什麼秘密了。關于
紅黑樹
的原理，請參考本人的另一篇文章-紅黑樹詳細分析，本篇文章不會對此展開讨論。
3.1 查找
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 循環裡 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; } 查找操作的核心邏輯就是
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 方法産生的集合，在編譯時會轉換成用疊代器周遊，等價于：
Set keys = map.keySet(); Iterator ite = keys.iterator(); while (ite.hasNext()) { Object key = ite.next(); // do something } 另一方面，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 的周遊操作就講到這。周遊操作本身不難，但講的有點多，略顯啰嗦，大家見怪。
3.3 插入
相對于前兩個操作，插入操作明顯要複雜一些。當往 TreeMap 中放入新的鍵值對後，可能會破壞紅黑樹的性質。這裡為了描述友善，把 Entry 稱為節點。并把新插入的節點稱為
N
，N 的父節點為
P
。P 的父節點為
G
，且 P 是 G 的左孩子。P 的兄弟節點為
U
。在往紅黑樹中插入新的節點 N 後（新節點為紅色），會産生下面5種情況：
- N 是根節點
- N 的父節點是黑色
- N 的父節點是紅色，叔叔節點也是紅色
- N 的父節點是紅色，叔叔節點是黑色，且 N 是 P 的右孩子
- N 的父節點是紅色，叔叔節點是黑色，且 N 是 P 的左孩子
上面5中情況中，情況2不會破壞紅黑樹性質，是以無需處理。情況1 會破壞紅黑樹性質2（根是黑色），情況3、4、和5會破壞紅黑樹性質4（每個紅色節點必須有兩個黑色的子節點）。這個時候就需要進行調整，以使紅黑樹重新恢複平衡。至于怎麼調整，可以參考我另一篇關于紅黑樹的文章（紅黑樹詳細分析），這裡不再重複說明。接下來分析一下插入操作相關源碼：
public V put(K key, V value) { Entry<K,V> t = root; // 1.如果根節點為 null，将新節點設為根節點 if (t == null) { compare(key, key); 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) { // 2.為 key 在紅黑樹找到合适的位置 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 { // 與上面代碼邏輯類似，省略 } Entry<K,V> e = new Entry<>(key, value, parent); // 3.将新節點鍊入紅黑樹中 if (cmp < 0) parent.left = e; else parent.right = e; // 4.插入新節點可能會破壞紅黑樹性質，這裡修正一下 fixAfterInsertion(e); size++; modCount++; return null; } 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 是
最終被删除
的節點，是原因二叉查找樹會把要删除有兩個孩子的節點的情況轉化為删除隻有一個孩子的節點的情況，該節點是欲被删除節點的前驅和後繼。
接下來，簡單列舉一下删除節點時可能會出現的情況，先列舉較為簡單的情況：
- 最終被删除的節點 X 是紅色節點
- X 是黑色節點，但該節點的孩子節點是紅色
比較複雜的情況：
- 替換節點 N 是新的根
- N 為黑色，N 的兄弟節點 S 為紅色，其他節點為黑色。
- N 為黑色，N 的父節點 P，兄弟節點 S 和 S 的孩子節點均為黑色。
- N 為黑色，P 是紅色，S 和 S 孩子均為黑色。
- N 為黑色，P 可紅可黑，S 為黑色，S 的左孩子 SL 為紅色，右孩子 SR 為黑色
- N 為黑色，P 可紅可黑，S 為黑色，SR 為紅色，SL 可紅可黑
上面列舉的8種情況中，前兩種處理起來比較簡單，後6種情況中情況26較為複雜。接下來我将會對情況26展開分析，删除相關的源碼如下：
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; } } } 從源碼中可以看出，
remove
方法隻是一個簡單的保證，核心實作在
deleteEntry
方法中。deleteEntry 主要做了這麼幾件事：
- 如果待删除節點 P 有兩個孩子，則先找到 P 的後繼 S，然後将 S 中的值拷貝到 P 中，并讓 P 指向 S
- 如果最終被删除節點 P（P 現在指向最終被删除節點）的孩子不為空，則用其孩子節點替換掉
- 如果最終被删除的節點是黑色的話，調用 fixAfterDeletion 方法進行修複
上面說了 replacement 不為空時，deleteEntry 的執行邏輯。上面說的略微啰嗦，如果簡單說的話，7個字即可總結：
找後繼 -> 替換 -> 修複
。這三步中，最複雜的是修複操作。修複操作要重新使紅黑樹恢複平衡，修複操作的源碼分析如下：
fixAfterDeletion 方法分析如下：
上面對 fixAfterDeletion 部分代碼邏輯就行了分析，通過配圖的形式解析了每段代碼邏輯所處理的情況。通過圖解，應該還是比較好了解的。好了，TreeMap 源碼先分析到這裡。
四、總結
本文可以看做是本人”紅黑樹詳細分析”一文的延續。前一篇文章從理論層面上詳細分析了紅黑樹插入和删除操作可能會導緻的問題，以及如何修複。本文則從實踐層面是分析了插入和删除操作在具體的實作中時怎樣做的。另外，本文選擇了從集合架構常用方法這一角度進行分析，詳細分析了查找、周遊、插入和删除等方法。總體來說，分析的還是比較詳細的。當然限于本人的水準，文中可能會存在一些錯誤的論述。如果大家發現了，歡迎指出來。如果這些錯誤的論述對你造成了困擾，我這裡先說聲抱歉。如果你也在學習 TreeMap 源碼，希望這篇文章能夠幫到你。
最後感謝大家花時間的閱讀我的文章，順祝大家寫代碼無BUG，下篇文章見。
本文在知識共享許可協定 4.0 下釋出，轉載請注明出處
作者：田小波
為了獲得更好的閱讀體驗，
請移步至本人的個人部落格：https://www.tianxiaobo.com
本作品采用知識共享署名-非商業性使用-禁止演繹 4.0 國際許可協定進行許可。
![Java小案例——随機數猜測随機數猜測[圖]](data:image/gif;base64,R0lGODlhAQABAIAAAP///wAAACwAAAAAAQABAAACAkQBADs=)