TreeMap源码解析
TreeMap是Map的一个实现类，底层数据结构是红黑树。它是有序的，默认顺序是自然排序，即ascii码升序，若在构造方法中传入比较器comparator，则按照比较器的规则进行排序。若对红黑树的数据结构不太熟悉，建议先移步https://blog.csdn.net/MAX_VALUE/article/details/104003631。本文主要介绍TreeMap的常用方法的工作流程，水平有限，若存在错误，请批评指正。
1.接口实现
NavigableMap
表示TreeMap是有序的，可拷贝的，可序列化的。
2.成员变量
//比较器
private final Comparator comparator;
//红黑树根节点
private transient Entry root;
//节点个数
private transient int size = 0;
//树结构被修改次数
private transient int modCount = 0;
//树节点
static final class Entry implements Map.Entry {
K key;
V value;
//左子树
Entry left;
//右子树
Entry right;
//父节点
Entry parent;
//节点颜色
boolean color = BLACK;
} 排序比较器comparator(默认按key的ascii码升序)，Entry
Entry就是一个红黑树节点，包含key,value,left,right,parent,color（left为左子树，ritht为右子树，parent为父节点，color为节点颜色）。
数据结构：红黑树
3.构造方法
//无参构造
public TreeMap() {
comparator = null;
}
//传入构造器
public TreeMap(Comparator comparator) {
this.comparator = comparator;
}
//传入Map其他实现类，相当于做转换
public TreeMap(Map m) {
comparator = null;
putAll(m);
}4.添加元素过程
public V put(K key, V value) {
Entry 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 parent;
// split comparator and comparable paths
Comparator 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();
@SuppressWarnings("unchecked")
Comparable k = (Comparable) key;
do {
//从根节点开始，比较和插入节点的key值大小
parent = t;
cmp = k.compareTo(t.key);
if (cmp < 0)
//将左子树设置为当前节点，继续查找
t = t.left;
else if (cmp > 0)
//将右子树设置为当前节点，继续查找
t = t.right;
else
//key值存在，更新value
return t.setValue(value);
} while (t != null);
}
Entry e = new Entry<>(key, value, parent);
if (cmp < 0)
//将新增节点设置为父节点的左子树
parent.left = e;
else
//将新增节点设置为父节点的右子树
parent.right = e;
//红黑树自平衡操作，基于红黑树插入场景算法
fixAfterInsertion(e);
//map中元素个数增加1，map结构被修改次数增加1
size++;
modCount++;
return null;
} /** From CLR */
//对照红黑树的插入场景图会比较好理解，以下代码就是操作场景的翻译过程
private void fixAfterInsertion(Entry x) {
//设置新增节点为红色节点
x.color = RED;
//插入节点的父节点是红色节点，破坏了红黑树的性质，需要做自平衡
while (x != null && x != root && x.parent.color == RED) {
//插入节点的父节点是祖父节点的左子节点
if (parentOf(x) == leftOf(parentOf(parentOf(x)))) {
//插入节点的叔叔节点（父节点的兄弟节点）
Entry y = rightOf(parentOf(parentOf(x)));
if (colorOf(y) == RED) {
//叔叔节点存在且为红色
//将父节点和叔叔节点设置为黑色，祖父节点设置为红色
setColor(parentOf(x), BLACK);
setColor(y, BLACK);
setColor(parentOf(parentOf(x)), RED);
//将祖父节点设置为插入节点x
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 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;
} 添加键值对，如果key值存在则替换旧的value值。
1.判断根节点是否存在，如果不存在将新增节点设置为根节点。结束
2.若根节点存在，判断是否传入比较器，如果有比较器根据比较器规则比较key值大小，如果未传入，则按照自然顺序进行排序（ascii码升序）。
3.找到插入节点的父节点。从根节点开始，比较插入节点的key值和当前节点的key值，如果大于将当前节点的右子树节点设置为当前节点，如果小于将当前节点的左子树节点设置为当前节点继续查找，如果等于则将该key值对应的value更新。循环判断直到当前节点为空，返回此时的父节点。
4.新增一个节点，该节点为红色。他的父节点是上述返回的父节点。
5.红黑树自平衡。场景具有对称性。3层判断（插入节点的父节点是否为红色、父节点是否是祖父节点的左子树、叔叔节点是否是红色、插入节点是父节点的左右子树）。
5.查找元素过程
根据key值查找指定的value
public V get(Object key) {
//根据key查找entry树节点
Entry p = getEntry(key);
return (p==null ? null : p.value);
}
final Entry getEntry(Object key) {
// Offload comparator-based version for sake of performance
if (comparator != null)
//构造treeMap时，传入了比较器，则按比较器的规则排序进行查询
//查找规则基于红黑树的查找算法
return getEntryUsingComparator(key);
if (key == null)
throw new NullPointerException();
@SuppressWarnings("unchecked")
Comparable k = (Comparable) key;
Entry p = root;
while (p != null) {
int cmp = k.compareTo(p.key);
if (cmp < 0)
//插入节点key小于当前节点key，向当前节点的左子树继续查找
p = p.left;
else if (cmp > 0)
//插入节点key大于当前节点key，向当前节点的右子树继续查找
p = p.right;
else
//插入节点key等于当前节点key，则p就是要查找的节点
return p;
}
return null;
} 1.传入比较器则按照比较器规则比较，否则按照自然顺序排序。
2.查找规则：将根节点设置为当前节点，若当前节点为空则直接返回null,若插入节点key值等于当前节点key值则返回当前节点，若小于则将当前节点的左子树设置为当前节点，若大于则将当前节点的右子树设置为当前节点。重复步骤2。
6.移除元素过程
public V remove(Object key) {
//查找要删除的树节点p
Entry p = getEntry(key);
if (p == null)
return null;
V oldValue = p.value;
//删除树节点p
deleteEntry(p);
return oldValue;
} private void deleteEntry(Entry p) {
//map结构变化次数增加1，节点数减1
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) {
//查找下一个节点，在此处应该是删除节点p的右子树的最左子节点，即替换节点
Entry s = successor(p);
//更新删除节点位置处的key,value为替换节点值
p.key = s.key;
p.value = s.value;
//p指向替换节点
p = s;
} // p has 2 children
//修正替换节点，删除节点只存在左子节点或只存在右子节点或不存在子节点
//1.若只存在左节点，则替换节点就是左子节点
//2.若只存在右节点，则替换节点就是右节点
//3.不存在子节点则替换节点为空
//4.删除节点的两个子节点都存在,若通过上述找到的替换节点存在右子节点，即替换节点
//不是叶子节点，此时将替换节点更改为它的右子节点。
//5.删除节点的两个子节点都存在,若通过上述找到的替换节点不存在右子节点,替代节点不变，
//replacement 为null。
// Start fixup at replacement node, if it exists.
Entry replacement = (p.left != null ? p.left : p.right);
if (replacement != null) {
// Link replacement to parent
//场景1、2、4
//解除删除节点引用方式1：使用替代节点替换删除节点的相关引用，删除节点的引用
//置为null。
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.
//场景3、5
if (p.color == BLACK)
fixAfterDeletion(p);
//解除删除节点引用方式2：删除节点的key,value值先替换成替换节点的key,value,
//然后解除替代节点的引用关系
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;
}
}
} /** From CLR */
//红色树删除操作场景算法
private void fixAfterDeletion(Entry x) {
//x为替换节点，替换节点为黑色节点，做自平衡
while (x != root && colorOf(x) == BLACK) {
//替换节点是父节点的左子节点
if (x == leftOf(parentOf(x))) {
//替换接地的兄弟节点
Entry 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 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);
} 根据key值移除键值对
1.根据key值查找到相应的节点，参考查找元素过程。
2.查找替代节点，若删除节点无子节点，则替代节点为本身；若删除节点有一个子节点，则替代节点为该子节点；若删除节点有两个子节点，则替代节点为右子树节点的最左节点。删除节点的方式有两种，方式1：使用替代节点替换删除节点的相关引用，删除节点的引用置为null。方式2：删除节点的key,value值先替换成替换节点的key,value,然后解除替代节点的引用关系。
3.红黑树完成自平衡。具有对称性。四层判断（替代节点是否为黑色，替代节点是否为其父节点的左子树，替代节点的兄弟节点是否为黑色，替代节点的兄弟节点的子树颜色）。确定好具体的场景进行变色和旋转，源码中的步骤同红黑树算法步骤。
7.TreeMap遍历
Map的遍历本质上还是基于Iterator的相关实现类。使用的是iterator()、next()方法。
for(Map.Entry entry : map.entrySet()) {} 遍历步骤：
public Set> entrySet() {
//第一次调用entrySet()时，entrySet为null。
EntrySet es = entrySet;
//new EntrySet()会调用其iterator()方法
return (es != null) ? es : (entrySet = new EntrySet());
}
class EntrySet extends AbstractSet> {
public Iterator> iterator() {
//查找红黑树的第一个节点，即最左子节点，遍历到的第一个元素就是该节点
return new EntryIterator(getFirstEntry());
}
}
public Map.Entry next() {
//之后调用iterator的next()方法
return nextEntry();
}
final Entry nextEntry() {
Entry e = next;
if (e == null)
throw new NoSuchElementException();
if (modCount != expectedModCount)
throw new ConcurrentModificationException();
//按照红黑树的排序顺序找到下一个节点
next = successor(e);
lastReturned = e;
return e;
} static TreeMap.Entry successor(Entry t) {
if (t == null)
return null;
else if (t.right != null) {
//当前节点的右子节点存在，则后继节点为右子节点或者是右子节点的最左子节点。
Entry p = t.right;
while (p.left != null)
p = p.left;
return p;
} else {
//1.若当前节点的是其父节点的右子节点，则下一个节点为null
//2.若当前节点的是其父节点的左子节点，则下一个节点是其父节点
Entry p = t.parent;
Entry ch = t;
while (p != null && ch == p.right) {
ch = p;
p = p.parent;
}
return p;
}
} 8.TreeSet源码分析
TreeSet
//m是TreeSet的底层数据结构
private transient NavigableMap m;
//通常使用无参构造
public TreeSet() {
this(new TreeMap());
}
//默认将TreeMap作为TreeSet底层的实现
TreeSet(NavigableMap m) {
this.m = m;
}