TreeMap
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 的情况
