Java集合源码分析（十四）：TreeMap
TreeMap是红黑树Java版的实现。如果还有小伙伴对红黑树还不是很了解的话，可以去之前的文章看一看。红黑树能保证对增删改查等基本操作的时间复杂度保持在O（lgN）.接下来就让我们一起来看看TreeMap的底层源码。
一、TreeMap继承结构
从图中我们可以看到，TreeMap继承了父类AbstractMap类，也实现了Cloneable、NavigableMap还有Serializable接口。这几个接口分别能带来什么这里就不一 一说明了。
二、Entry类
static final class Entry<K,V> implements Map.Entry<K,V> {
K key;
V value;
Entry<K,V> left;
Entry<K,V> right;
Entry<K,V> parent;
boolean color = BLACK;
/**
* Make a new cell with given key, value, and parent, and with
* {@code null} child links, and BLACK color.
*/
Entry(K key, V value, Entry<K,V> parent) {
this.key = key;
this.value = value;
this.parent = parent;
}
/**
* Returns the key.
*
* @return the key
*/
public K getKey() {
return key;
}
/**
* Returns the value associated with the key.
*
* @return the value associated with the key
*/
public V getValue() {
return value;
}
/**
* Replaces the value currently associated with the key with the given
* value.
*
* @return the value associated with the key before this method was
* called
*/
public V setValue(V value) {
V oldValue = this.value;
this.value = value;
return oldValue;
}
public boolean equals(Object o) {
if (!(o instanceof Map.Entry))
return false;
Map.Entry<?,?> e = (Map.Entry<?,?>)o;
return valEquals(key,e.getKey()) && valEquals(value,e.getValue());
}
public int hashCode() {
int keyHash = (key==null ? 0 : key.hashCode());
int valueHash = (value==null ? 0 : value.hashCode());
return keyHash ^ valueHash;
}
public String toString() {
return key + "=" + value;
}
}
可以看到TreeMap里的Entry类实现了Map.Entry接口，以一个Entry为一个节点，分别有左孩子，右孩子和父节点，默认的颜色为黑色。
二、构造函数
// 默认构造，比较器采用key的自然比较顺序
public TreeMap() {
comparator = null;
}
// 指定比较器
public TreeMap(Comparator<? super K> comparator) {
this.comparator = comparator;
}
// 从Map集合导入初始数据
public TreeMap(Map<? extends K, ? extends V> m) {
comparator = null;
putAll(m);
}
// 从SortedMap导入初始数据
public TreeMap(SortedMap<K, ? extends V> m) {
comparator = m.comparator();
try {
buildFromSorted(m.size(), m.entrySet().iterator(), null, null);
} catch (java.io.IOException cannotHappen) {
} catch (ClassNotFoundException cannotHappen) {
}
}
前两个构造器我想不用多说，后面两个构造器分别使用了putAll()方法和 buildFromSorted（）方法，后面这两个方法我们会慢慢分析，现在继续其他的内容。
三、TreeMap一些重要的实现方法
红黑树最麻烦也是最难的部分就是增和删了，我们从增加一个元素开始分析起走：
来看看源码：
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;
// 初始化时指定了comparator比较器
if (cpr != null) {
do {
// 把t暂存到parent中
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 {
// 使用key的比较器，while循环原理和上述一致
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);
}
// 不断的比较，找到了没有相应儿子的节点
//（cmp<0就是没有左儿子，cmp>0就是没有右儿子）
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;
}
过程分析：首先暂存一个当前的根节点。然后进行判断。如果当前的根节点为空，也就是还没有元素的时候，然后在执行compare方法比较Key值，（个人觉得这句代码没有任何意义，空树还需要比较、排序？）
final int compare(Object k1, Object k2) {
return comparator==null ? ((Comparable<? super K>)k1).compareTo((K)k2)
: comparator.compare((K)k1, (K)k2);
}
接下来再将新的key-value键值对装进Entry，并且设置为根节点，然后size变为1，修改次数+1； 如果根节点不为空的话，这里定义了一个int型的cmp变量，这里也不知道这个变量是干嘛的，接着往后面看，指定的排序算法comparator,如果cpr不为null的话，就按照上一步的排序算法进行创建TreeMap集合，是怎么创建的呢？
将parent节点指向上次循环后的t节点，然后比较新增节点的key和当前节点的key，如果cmp返回值小于0，表示新增节点的key小于当前节点的key，则以当前节点的左子节点作为新的当前节点，如果cmp返回值大于0，表示新增节点的key大于当前节点的key，则以当前节点的右子节点作为新的当前节点，最后cmp返回值等于0，表示两个key值相等，则新值覆盖旧值，并返回新值；如果cpr为null，则采用默认的排序算法进行创建TreeMap集合，key值为空的话。抛出异常，然后处理过程就和上面cpr不为空时是一样的。
通过上述步骤确定了父节点后，就将新增的这个节点根据值的大小放到当前父节点的右左边，然后size+1，修改次数+1，最重要的一个难点来了，我们看到，在把数据插入后，执行了一个 fixAfterInsertion(e)方法，新插入的元素破坏了红黑树规则，需要调整，
我们来看看这个功能是在怎么实现的：
private void fixAfterInsertion(Entry<K,V> x) {
// 先把x节点染成红色，这样可以不增加黑高，简化调整问题
x.color = RED;
// 条件是父节点是红色的，且x不是root节点，
// 因为到root节点后就走到另外的分支了，而那个分支是正确的
while (x != null && x != root && x.parent.color == RED) {
//x的父节点是其祖父节点的左儿子
if (parentOf(x) == leftOf(parentOf(parentOf(x)))) {
// y是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是右节点，以其父节点左旋
x = parentOf(x);
rotateLeft(x);
}
// 右旋
setColor(parentOf(x), BLACK);
setColor(parentOf(parentOf(x)), RED);
rotateRight(parentOf(parentOf(x)));
}
} else {
//x的父节点是其祖父节点的右儿子
// y是其叔叔
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是左节点，以其父节点右旋
x = parentOf(x);
rotateRight(x);
}
//左旋
setColor(parentOf(x), BLACK);
setColor(parentOf(parentOf(x)), RED);
rotateLeft(parentOf(parentOf(x)));
}
}
}
//root节点颜色为黑色
root.color = BLACK;
}
分析过程：先把新增的节点的颜色变为红色，这样就不存在黑高问题（每条路径上的黑色节点数必须一样），接下俩就进入到一个循环了，循环终止的条件为当前父节点的颜色为红色、当前节点不为空并且当前节点不是根节点，循环里又是怎么做的呢？
这里请参照上面的代码和注释，后续我们会通过图解的方式来说明这个过程，上面的代码也涉及到左旋和右旋代码：
// 右旋与左旋思路一致，只分析其一
// 结果相当于把p和p的儿子调换了
private void rotateLeft(Entry<K,V> p) {
if (p != null) {
// 取出p的右儿子
Entry<K,V> r = p.right;
// 然后将p的右儿子的左儿子，也就是p的左孙子变成p的右儿子
p.right = r.left;
if (r.left != null)
// p的左孙子的父亲现在是p
r.left.parent = p;
// 然后把p的父亲，设置为p右儿子的父亲
r.parent = p.parent;
// 这说明p原来是root节点
if (p.parent == null)
root = r;
else if (p.parent.left == p)
p.parent.left = r;
else
p.parent.right = r;
r.left = p;
p.parent = r;
}
}
//和左旋类似
private void rotateRight(Entry<K,V> p) {
// ...
}
在图解之前，我们的熟悉红黑树的定义，之前已说明过，这里在贴出来：
1、 每个节点或者是黑色，或者是红色。
2、 根节点是黑色。
3、 每个叶子节点（NIL）是黑色。 [注意：这里叶子节点，是指为空(NIL或NULL)的叶子节点！
4、 如果一个节点是红色的，则它的两个儿子都是黑色的。
5、 从一个节点到该节点的子孙节点的所有路径上包含相同数目的黑节点。
假如现在有一颗红黑树，把元素7插入到下面的红黑树中，：
我们根据左子树>父节点>右节点，一个一个的比较，直到发现8是一个叶节点，所以把7插入到8的左儿子处，也就是下面这样：
为了不增加黑高，这里把7设置为红色。现在，这棵树已经不再是红黑树了，因为其违反了规则如果一个节点是红色的，则它的两个儿子都是黑色的。我们按照fixAfterInsertion的方式对其进行调整，fixAfterInsertion中的参数x就是这里的7。
我们来看看调整过程。
首先，进入循环后发现的7的父亲是右节点，进入else判断，7的叔叔为4是红色的，于是把4和8染成黑色，6染为红色，把x参数指向6，并进入下一次循环：
此时此刻，x是6，其父亲10是左儿子，其叔叔18是黑色的，此时代码会走到这一步：
if (x == rightOf(parentOf(x))) {
x = parentOf(x);
rotateLeft(x);
}
setColor(parentOf(x), BLACK);
setColor(parentOf(parentOf(x)), RED);
rotateRight(parentOf(parentOf(x)));
也就是说把10和14的颜色更换，如下图所示：
然后以14为基础右旋，涉及到的元素有10、12和14，如下图所示：
具体操作为，把10的右儿子12，变为14的左儿子，然后再把14变为10的右儿子：
循环到这一步也就跳出截止了，我么再来检查检查，这是一颗正确的没毛病的红黑树。
还会有一些其他更复杂的情况，大家可以参考之前的文章 红黄树Java集合源码分析（六）：红黑树
再来看看获取元素
获取元素
TreeMap中的元素是有序的，当使用中序遍历时就可以得到一个有序的Set集合，所以获取元素可以采用二分法：
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 (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;
}
除了获取某个元素外，还可以获取它的前一个或者后一个元素：
// 获取前一个元素
static <K,V> Entry<K,V> predecessor(Entry<K,V> t) {
if (t == null)
return null;
else if (t.left != null) {
// t有左孩子，所以t的前一个元素是它左孩子所在的子树的最右侧叶子结点
Entry<K,V> p = t.left;
while (p.right != null)
p = p.right;
return p;
} else {
// t没有左孩子，所以t的前一个元素有两种情况
// 1. t是右孩子，那它的前一个元素就是它的父结点
// 2. t是左孩子，它的前一个元素需要向上递归，直到递归到下一个是右孩子的节点，转为情况1
Entry<K,V> p = t.parent;
Entry<K,V> ch = t;
while (p != null && ch == p.left) {
ch = p;
p = p.parent;
}
return p;
}
}
// 获取后一个元素
static <K,V> TreeMap.Entry<K,V> successor(Entry<K,V> t) {
//...
}
删除一个元素
我们看看java代码是怎么实现的：
public V remove(Object key) {
// 先用二分法获取这个元素，如果为null，不需要继续了
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--;
// If strictly internal, copy successor's element to p and then make p
// point to successor.
//如果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.
// p有一个儿子，或者没有儿子，获取到之后放在replacement中
Entry<K,V> replacement = (p.left != null ? p.left : p.right);
// p有儿子
if (replacement != null) {
// Link replacement to parent
// 把p的子孙接在p的父级
replacement.parent = p.parent;
//p是根节点
if (p.parent == null)
root = replacement;
//p是左儿子
else if (p == p.parent.left)
p.parent.left = replacement;
// p是右儿子
else
p.parent.right = replacement;
//把p的链接都删掉
// 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 {
//p没有儿子
// No children. Use self as phantom replacement and unlink.
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;
}
}
}
修正调整的方法：
private void fixAfterDeletion(Entry<K,V> x) {
while (x != root && colorOf(x) == BLACK) {
// x是左儿子
if (x == leftOf(parentOf(x))) {
// sib是x的兄弟
Entry<K,V> sib = rightOf(parentOf(x));
// 兄弟是红色的
if (colorOf(sib) == RED) {
setColor(sib, BLACK);
setColor(parentOf(x), RED);
rotateLeft(parentOf(x));
sib = rightOf(parentOf(x));
}
// 兄弟没有孩子或者孩子是黑色的
if (colorOf(leftOf(sib)) == BLACK &&
colorOf(rightOf(sib)) == BLACK) {
setColor(sib, RED);
x = parentOf(x);
} else {
// 兄弟的右孩子是黑色的
if (colorOf(rightOf(sib)) == BLACK) {
setColor(leftOf(sib), BLACK);
setColor(sib, RED);
rotateRight(sib);
sib = rightOf(parentOf(x));
}
setColor(sib, colorOf(parentOf(x)));
setColor(parentOf(x), BLACK);
setColor(rightOf(sib), BLACK);
rotateLeft(parentOf(x));
x = root;
}
} else {
// symmetric
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);
}
删除元素的图示过程在之前的文章中也有讲到，这里就不在一 一说明了。
其他重要方法
在前面分析构造函数时，有两个函数putAll和buildFromSorted当时忽略了，现在我们来看看它们的实现。
public void putAll(Map<? extends K, ? extends V> map) {
int mapSize = map.size();
if (size==0 && mapSize!=0 && map instanceof SortedMap) {
//...
buildFromSorted(mapSize, map.entrySet().iterator(),
null, null);
//...
return;
}
super.putAll(map);
}
可以看出来，putAll方法，当Map是一个SortedMap实例时，依赖于buildFromSorted，其他情况则是由AbstractMap实现的。所以这里重点看下buildFromSorted的实现。
buildFromSorted有两个，一个是供putAll等调用的，另外一个则是具体的实现。
// 这个方法主要是被调用，关注它只为了看下computeRedLevel这个方法
private void buildFromSorted(int size, Iterator<?> it,
java.io.ObjectInputStream str,
V defaultVal)
throws java.io.IOException, ClassNotFoundException {
this.size = size;
root = buildFromSorted(0, 0, size - 1, computeRedLevel(size),
it, str, defaultVal);
}
这里调用了一个computeRedLevel的方法，是这里的关键。
private static int computeRedLevel(int sz) {
int level = 0;
for (int m = sz - 1; m >= 0; m = m / 2 - 1)
level++;
return level;
}
这个方法和染色为红色有关，其实现和二分法看似有一定联系，通过这种方式可以构建一个优秀的红黑树，能够为以后插入更多数据提供便利。
最后我们看下buildFromSorted的实现：
private final Entry<K,V> buildFromSorted(int level, int lo, int hi,
int redLevel,
Iterator<?> it,
java.io.ObjectInputStream str,
V defaultVal)
throws java.io.IOException, ClassNotFoundException {
if (hi < lo) return null;
// 获取中间位置
int mid = (lo + hi) >>> 1;
Entry<K,V> left = null;
if (lo < mid)
// 递归左子树，和压栈类似，直到lo>=mid才能返回结果
left = buildFromSorted(level+1, lo, mid - 1, redLevel,
it, str, defaultVal);
// extract key and/or value from iterator or stream
K key;
V value;
if (it != null) {
// 给key和value赋值
if (defaultVal==null) {
Map.Entry<?,?> entry = (Map.Entry<?,?>)it.next();
key = (K)entry.getKey();
value = (V)entry.getValue();
} else {
key = (K)it.next();
value = defaultVal;
}
} else {
// use stream
// 从序列化中恢复
key = (K) str.readObject();
value = (defaultVal != null ? defaultVal : (V) str.readObject());
}
Entry<K,V> middle = new Entry<>(key, value, null);
// color nodes in non-full bottommost level red
//
if (level == redLevel)
middle.color = RED;
if (left != null) {
middle.left = left;
left.parent = middle;
}
if (mid < hi) {
Entry<K,V> right = buildFromSorted(level+1, mid+1, hi, redLevel,
it, str, defaultVal);
middle.right = right;
right.parent = middle;
}
return middle;
}
举例插入10条数据，得到以下结果：
可以看出来，redLevel控制的是红色节点出现的层级，使插入的数据更整齐，方便后续操作。
这次的讨论就到这里，欢迎留言讨论！