【集合系列】- 深入淺出的分析TreeMap
一、摘要
在集合系列的第一章，咱們瞭解到，Map的實現類有HashMap、LinkedHashMap、TreeMap、IdentityHashMap、WeakHashMap、Hashtable、Properties等等。
本文主要從資料結構和演算法層面，探討TreeMap的實現。
二、簡介
Java TreeMap實現了SortedMap介面，也就是說會按照key的大小順序對Map中的元素進行排序，key大小的評判可以通過其本身的自然順序（natural ordering），也可以通過構造時傳入的比較器（Comparator）。
TreeMap底層通過紅黑樹（Red-Black tree）實現，所以要了解TreeMap就必須對紅黑樹有一定的瞭解，在《集合系列》文章中，如果你已經讀過紅黑樹的講解，其實本文要講解的TreeMap，跟其大同小異。
紅黑樹又稱紅-黑二叉樹，它首先是一顆二叉樹，它具有二叉樹所有的特性。同時紅黑樹更是一顆自平衡的排序二叉樹。
對於一棵有效的紅黑樹二叉樹，主要有以下規則：
- 1、每個節點要麼是紅色，要麼是黑色，但根節點永遠是黑色的；
- 2、每個紅色節點的兩個子節點一定都是黑色；
- 3、紅色節點不能連續（也即是，紅色節點的孩子和父親都不能是紅色）；
- 4、從任一節點到其子樹中每個葉子節點的路徑都包含相同數量的黑色節點；
- 5、所有的葉節點都是是黑色的（注意這裡說葉子節點其實是上圖中的 NIL 節點）；
這些約束強制了紅黑樹的關鍵性質：從根到葉子的最長的可能路徑不多於最短的可能路徑的兩倍長。結果是這棵樹大致上是平衡的。因為操作比如插入、刪除和查詢某個值的最壞情況時間都要求與樹的高度成比例，這個在高度上的理論上限允許紅黑樹在最壞情況下都是高效的，而不同於普通的二叉查詢樹。所以紅黑樹它是複雜而高效的，其檢索效率O(log n)。下圖為一顆典型的紅黑二叉樹。
在樹的結構發生改變時（插入或者刪除操作），往往會破壞上述條件3或條件4，需要通過調整使得查詢樹重新滿足紅黑樹的條件。
調整方式主要有：左旋、右旋和顏色轉換！
2.1、左旋
左旋的過程是將x的右子樹繞x逆時針旋轉，使得x的右子樹成為x的父親，同時修改相關節點的引用。旋轉之後，二叉查詢樹的屬性仍然滿足。
2.2、右旋
右旋的過程是將x的左子樹繞x順時針旋轉，使得x的左子樹成為x的父親，同時修改相關節點的引用。旋轉之後，二叉查詢樹的屬性仍然滿足。
2.3、顏色轉換
顏色轉換的過程是將紅色節點轉換為黑色節點，或者將黑色節點轉換為紅色節點，以滿足紅黑樹二叉樹的規則！
三、常用方法介紹
3.1、get方法
get方法根據指定的key值返回對應的value，該方法呼叫了getEntry(Object key)得到相應的entry，然後返回entry.value。
演算法思想是根據key的自然順序（或者比較器順序）對二叉查詢樹進行查詢，直到找到滿足k.compareTo(p.key) == 0的entry。
原始碼如下：
final Entry<K,V> getEntry(Object key) {
//如果傳入比較器，通過getEntryUsingComparator方法獲取元素
if (comparator != null)
return getEntryUsingComparator(key);
//不允許key值為null
if (key == null)
throw new NullPointerException();
//使用元素的自然順序
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;
}如果傳入比較器，通過getEntryUsingComparator方法獲取元素
final Entry<K,V> getEntryUsingComparator(Object key) {
K k = (K) key;
Comparator<? super K> cpr = comparator;
if (cpr != null) {
Entry<K,V> p = root;
while (p != null) {
//通過比較器順序，獲取元素
int cmp = cpr.compare(k, p.key);
if (cmp < 0)
p = p.left;
else if (cmp > 0)
p = p.right;
else
return p;
}
}
return null;
}測試用例：
public static void main(String[] args) {
Map initMap = new TreeMap();
initMap.put("4", "d");
initMap.put("3", "c");
initMap.put("1", "a");
initMap.put("2", "b");
//預設自然排序，key為升序
System.out.println("預設 排序結果:" + initMap.toString());
//自定義排序，在TreeMap初始化階段傳入Comparator 內部物件
Map comparatorMap = new TreeMap<String, String>(new Comparator<String>() {
@Override
public int compare(String o1, String o2){
//根據key比較大小，採用倒敘，以大到小排序
return o2.compareTo(o1);
}
});
comparatorMap.put("4", "d");
comparatorMap.put("3", "c");
comparatorMap.put("1", "a");
comparatorMap.put("2", "b");
System.out.println("自定義 排序結果:" + comparatorMap.toString());
}輸出結果：
預設 排序結果:{1=a, 2=b, 3=c, 4=d}
自定義 排序結果:{4=d, 3=c, 2=b, 1=a}3.2、put方法
put方法是將指定的key, value對新增到map裡。該方法首先會對map做一次查詢，看是否包含該元組，如果已經包含則直接返回，查詢過程類似於getEntry()方法；如果沒有找到則會在紅黑樹中插入新的entry，如果插入之後破壞了紅黑樹的約束，還需要進行調整（旋轉，改變某些節點的顏色）。
原始碼如下：
public V put(K key, V value) {
Entry<K,V> t = root;
//如果紅黑樹根部為空，直接插入
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) {
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();
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
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;
}紅黑樹調整函式fixAfterInsertion(Entry<K,V> x)
private void fixAfterInsertion(Entry<K,V> x) {
x.color = RED;
while (x != null && x != root && x.parent.color == RED) {
//判斷x是否在樹的左邊，還是右邊
if (parentOf(x) == leftOf(parentOf(parentOf(x)))) {
Entry<K,V> y = rightOf(parentOf(parentOf(x)));
//如果x的父親的父親的右子樹是紅色，違反了紅色節點不能連續
if (colorOf(y) == RED) {
//進行顏色調整，以滿足紅色節點不能連續規則
setColor(parentOf(x), BLACK);
setColor(y, BLACK);
setColor(parentOf(parentOf(x)), RED);
x = parentOf(parentOf(x));
} else {
//如果x的父親的右子樹等於自己，那麼需要進行左旋轉
if (x == rightOf(parentOf(x))) {
x = parentOf(x);
rotateLeft(x);
}
setColor(parentOf(x), BLACK);
setColor(parentOf(parentOf(x)), RED);
rotateRight(parentOf(parentOf(x)));
}
} else {
//跟上面的流程正好相反
//獲取x的父親的父親的左子樹節點
Entry<K,V> y = leftOf(parentOf(parentOf(x)));
//如果y是紅色節點，違反了紅色節點不能連續
if (colorOf(y) == RED) {
//進行顏色轉換
setColor(parentOf(x), BLACK);
setColor(y, BLACK);
setColor(parentOf(parentOf(x)), RED);
x = parentOf(parentOf(x));
} else {
//如果x的父親的左子樹等於自己，那麼需要進行右旋轉
if (x == leftOf(parentOf(x))) {
x = parentOf(x);
rotateRight(x);
}
setColor(parentOf(x), BLACK);
setColor(parentOf(parentOf(x)), RED);
rotateLeft(parentOf(parentOf(x)));
}
}
}
//根節點一定為黑色
root.color = BLACK;
}上述程式碼的插入流程：
- 1、首先在紅黑樹上找到合適的位置；
- 2、然後建立新的entry並插入；
- 3、通過函式fixAfterInsertion()，對某些節點進行旋轉、改變某些節點的顏色，進行調整；
調整圖解：
3.3、remove方法
remove的作用是刪除key值對應的entry，該方法首先通過上文中提到的getEntry(Object key)方法找到 key 值對應的 entry，然後呼叫deleteEntry(Entry<K,V> entry)刪除對應的 entry。由於刪除操作會改變紅黑樹的結構，有可能破壞紅黑樹的約束，因此有可能要進行調整。
原始碼如下：
public V remove(Object key) {
Entry<K,V> p = getEntry(key);
if (p == null)
return null;
V oldValue = p.value;
deleteEntry(p);
return oldValue;
}刪除函式 deleteEntry()
private void deleteEntry(Entry<K,V> p) {
modCount++;
size--;
if (p.left != null && p.right != null) {// 刪除點p的左右子樹都非空。
Entry<K,V> s = successor(p);// 後繼
p.key = s.key;
p.value = s.value;
p = s;
}
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;
p.left = p.right = p.parent = null;
if (p.color == BLACK)
fixAfterDeletion(replacement);// 調整
} else if (p.parent == null) {
root = null;
} else { //刪除點p的左右子樹都為空
if (p.color == BLACK)
fixAfterDeletion(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;
}
}
}刪除後調整函式fixAfterDeletion()的具體程式碼如下：
private void fixAfterDeletion(Entry<K,V> x) {
while (x != root && colorOf(x) == BLACK) {
//判斷當前刪除的元素，是在x父親的左邊還是右邊
if (x == leftOf(parentOf(x))) {
Entry<K,V> sib = rightOf(parentOf(x));
//判斷x的父親的右子樹，是紅色還是黑色節點
if (colorOf(sib) == RED) {
//進行顏色轉換
setColor(sib, BLACK);
setColor(parentOf(x), RED);
rotateLeft(parentOf(x));
sib = rightOf(parentOf(x));
}
//x的父親的右子樹的左邊是黑色節點，右邊也是黑色節點
if (colorOf(leftOf(sib)) == BLACK &&
colorOf(rightOf(sib)) == BLACK) {
//設定x的父親的右子樹為紅色節點，將x的父親賦值給x
setColor(sib, RED);
x = parentOf(x);
} else {
//x的父親的右子樹的左邊是紅色節點，右邊也是黑色節點
if (colorOf(rightOf(sib)) == BLACK) {
//x的父親的右子樹的左邊進行顏色調整，右旋調整
setColor(leftOf(sib), BLACK);
setColor(sib, RED);
rotateRight(sib);
sib = rightOf(parentOf(x));
}
//對x進行左旋，顏色調整
setColor(sib, colorOf(parentOf(x)));
setColor(parentOf(x), BLACK);
setColor(rightOf(sib), BLACK);
rotateLeft(parentOf(x));
x = root;
}
} else { // 跟前四種情況對稱
Entry<K,V> sib = leftOf(parentOf(x));
if (colorOf(sib) == RED) {
setColor(sib, BLACK);
setColor(parentOf(x), RED);
rotateRight(parentOf(x));
sib = leftOf(parentOf(x));
}
if (colorOf(rightOf(sib)) == BLACK &&
colorOf(leftOf(sib)) == BLACK) {
setColor(sib, RED);
x = parentOf(x);
} else {
if (colorOf(leftOf(sib)) == BLACK) {
setColor(rightOf(sib), BLACK);
setColor(sib, RED);
rotateLeft(sib);
sib = leftOf(parentOf(x));
}
setColor(sib, colorOf(parentOf(x)));
setColor(parentOf(x), BLACK);
setColor(leftOf(sib), BLACK);
rotateRight(parentOf(x));
x = root;
}
}
}
setColor(x, BLACK);
}上述程式碼的刪除流程：
- 1、首先在紅黑樹上找到合適的位置；
- 2、然後刪除entry；
- 3、通過函式fixAfterDeletion()，對某些節點進行旋轉、改變某些節點的顏色，進行調整；
四、總結
TreeMap 預設是按鍵值的升序排序，如果需要自定義排序，可以通過new Comparator構造引數，重寫compare方法，進行自定義比較。
以上，主要是對 java 集合中的 TreeMap 做了寫講解，如果有理解不當之處，歡迎指正。
五、參考
1、JDK1.7&JDK1.8 原始碼
2、部落格園 - chenssy - TreeMap分析
2、知乎 - CarpenterLee - TreeMap講