Collection 系列文章的總目錄:
- Collection 體系的三個核心約定
- Sorted & Navigable
- Iterator & Iterable
- Java 中的數組
- ArrayList
- LinkedList
- HashMap
- LinkedHashMap
- TreeMap
- HashSet/LinkedHashSet/TreeSet
TreeMap 是一個有序的 Map，它能保證 key 的大小有序，內部使用紅黑樹來實現。
原理：
比較器：
要實現大小有序，就需要對比兩個元素的大小，TreeMap 支持兩種比較器：
- key 自身實現了 Comparable 接口，TreeMap 會調用 compareTo() 進行比較
- 傳入一個 Comparator，TreeMap 會調用 compare() 將兩個元素傳入，交給 comparator 比較
添加元素：
put(K key, V value)：
- 如果是第一次 put，直接創建一個節點，並賦值給 root，並返回
- 如果 comparator 不為空，則使用它的 compare() 進行比較
- 否則，將 key 強轉成 Comparable，然后使用 key 的 compareTo() 進行比較
- 如果 put 的 key 比當前節點的 key 小， 則比較左子樹
- 如果比當前的 key 大，則比較右子樹
- 如果相等，則替換 value，然后返回被替換的值
- 如果沒找到相等的節點，parent 會指向最后一個遍歷的葉子節點
- 創建一個新節點，和 parent 的 key 比較
- 如果小於 parent 的 key，則掛在左子樹
- 如果大於 parent 的 key，則掛在右子樹
- 調用 fixAfterInsertion() 通過旋轉紅黑樹，使之平衡
public V put(K key, V value) {
Entry<K,V> t = root;
// 如果第一次put
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;
// 如果比較器不為空，則用比較器進行比較。否則調用key的compareTo進行比較
if (cpr != null) {
// 注意這里沒有對key進行null檢查，說明在傳入比較器的時候，你可以put一個null的key
do {
parent = t;
cmp = cpr.compare(key, t.key);
// 如果put的key小於當前節點的key，對比左子樹
// 如果大於，對比右子樹
// 如果等於，覆蓋值，setValue
if (cmp < 0)
t = t.left;
else if (cmp > 0)
t = t.right;
else
return t.setValue(value);
} while (t != null);
// 當parent = t為null的時候，說明插入的是新節點，
// 並且parent指向了准備插入的節點的父節點，簡稱葉子父節點
}
else {
// 使用key的compareTo的流程也一樣
// 提前進行null檢查，因為下面會將key強轉成Comparable，並調用compareTo()
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;
}
// 如果comparator為null，則當k1為null的時候，會拋出NullPointerException
final int compare(Object k1, Object k2) {
return comparator==null ? ((Comparable<? super K>)k1).compareTo((K)k2)
: comparator.compare((K)k1, (K)k2);
}
插入節點后紅黑樹的旋轉：
// CLR是算法導論中三個作者名字的縮寫
/** From CLR */
private void fixAfterInsertion(Entry<K,V> x) {
x.color = RED;
while (x != null && x != root && x.parent.color == RED) {
if (parentOf(x) == leftOf(parentOf(parentOf(x)))) {
Entry<K,V> y = rightOf(parentOf(parentOf(x)));
if (colorOf(y) == RED) {
setColor(parentOf(x), BLACK);
setColor(y, BLACK);
setColor(parentOf(parentOf(x)), RED);
x = parentOf(parentOf(x));
} else {
if (x == rightOf(parentOf(x))) {
x = parentOf(x);
rotateLeft(x);
}
setColor(parentOf(x), BLACK);
setColor(parentOf(parentOf(x)), RED);
rotateRight(parentOf(parentOf(x)));
}
} else {
Entry<K,V> y = leftOf(parentOf(parentOf(x)));
if (colorOf(y) == RED) {
setColor(parentOf(x), BLACK);
setColor(y, BLACK);
setColor(parentOf(parentOf(x)), RED);
x = parentOf(parentOf(x));
} else {
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;
}
刪除元素：
remove(Object key)：
- 根據 key 找到對應的節點
- 使用傳入的 comparator 的 compare() 或 key 的 compareTo() 進行比較
- 小於，則比較左子樹
- 大於，則比較右子樹
- 等於，則返回當前節點
- 沒找到，返回 null
- 將節點的 value 保存到局部變量
- 刪除節點，並通過旋轉紅黑樹，使之平衡
- 返回節點的 value
public V remove(Object key) {
// 尋找元素
Entry<K,V> p = getEntry(key);
if (p == null)
return null;
V oldValue = p.value;
// 刪除元素
deleteEntry(p);
return oldValue;
}
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;
// 和前面的put差不多，就是普通的二叉查找樹的流程
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;
}
// 如果沒有找到，返回null
return null;
}
// 感覺兩部分代碼可以抽取共用，但是不知道為什么JDK沒有這么做
// 使用comaprator獲取元素
final Entry<K,V> getEntryUsingComparator(Object key) {
@SuppressWarnings("unchecked")
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;
}
刪除節點以及紅黑樹的旋轉：
private void deleteEntry(Entry<K,V> p) {
modCount++;
size--;
// If strictly internal, copy successor's element to p and then make p
// point to successor.
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) {
// Link replacement to parent
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;
// Fix replacement
if (p.color == BLACK)
fixAfterDeletion(replacement);
} else if (p.parent == null) { // return if we are the only node.
root = null;
} else { // No children. Use self as phantom replacement and unlink.
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;
}
}
}
問題：
TreeMap 的 key 不能為 null 嗎？
- 如果不傳入 Comparator，則 key 是不能為 null
- 如果傳入 Comparator，則有 Comparator 的實現決定
所以在使用 Comparator 的時候，靠注意 null 的情況