HashMap的红黑树实现源码分析
PS: 最近看了jdk的TreeMap、HashMap的红黑树代码，就动手用java实现了二叉树的数据结构，做了泛型封装，代码有注释，下载地址：
红黑树、二叉平衡树、二叉排序树的java实现
效果大致如下：
相关文章：HashMap源码分析
一、链表转红黑树
HashMap有两个成员变量TREEIFY_THRESHOLD、MIN_TREEIFY_CAPACITY。
- 当链表长度达到TREEIFY_THRESHOLD-1，就会检查是否扩容还是把链表结构转红黑树结构；
- 而当数组的长度小于MIN_TREEIFY_CAPACITY就会对数字进行扩容，否则就把链表结构转红黑树结构
从treeifyBin()方法看起
final void treeifyBin(Node<K,V>[] tab, int hash) {
int n, index; Node<K,V> e;
if (tab == null || (n = tab.length) < MIN_TREEIFY_CAPACITY)
resize();//开始扩容
else if ((e = tab[index = (n - 1) & hash]) != null) {//转红黑树结构
TreeNode<K,V> hd = null, tl = null;
do {//【标记1】把链表节点都改成红黑树节点
TreeNode<K,V> p = replacementTreeNode(e, null);
if (tl == null)
hd = p;
else {
p.prev = tl;
tl.next = p;
}
tl = p;
} while ((e = e.next) != null);
//这时还是一条链表，不过节点都变成了树节点
if ((tab[index] = hd) != null)
hd.treeify(tab);//【标记2】开始构建树
}
}
【标记1】链表节点改成红黑树节点很简单
TreeNode<K,V> replacementTreeNode(Node<K,V> p, Node<K,V> next) {
return new TreeNode<>(p.hash, p.key, p.value, next);
}
static final class TreeNode<K,V> extends LinkedHashMap.Entry<K,V> {
TreeNode<K,V> parent;
TreeNode<K,V> left;
TreeNode<K,V> right;
TreeNode<K,V> prev;//删除的时候有用
boolean red;
TreeNode(int hash, K key, V val, Node<K,V> next) {
super(hash, key, val, next);
}
...
}
TreeNode继承了LinkedHashMap.Entry，新增了父节点、左右节点、前节点和是否为红色节点的成员变量
【标记2】接下来是构建树treeify()方法。（该方法是TreeNode中的成员方法，下面的this指的是TreeNode实例）
final void treeify(Node<K,V>[] tab) {
TreeNode<K,V> root = null;//红黑树根节点
for (TreeNode<K,V> x = this, next; x != null; x = next) {
next = (TreeNode<K,V>)x.next;
x.left = x.right = null;
if (root == null) {//一开始为null
x.parent = null;
x.red = false;//红黑树根节点必须是黑色
root = x;//直接赋值给根节点
}
else {//这里说明根节点已存在，就要按二叉排序树方式插入
K k = x.key;
int h = x.hash;
Class<?> kc = null;
for (TreeNode<K,V> p = root;;) {//从根节点开始找合适插入位置
int dir, ph;
K pk = p.key;
if ((ph = p.hash) > h)
dir = -1;//往左子树走
else if (ph < h)
dir = 1;//往右子树走
//hash值相等，就进一步处理
else if ((kc == null &&
(kc = comparableClassFor(k)) == null) ||
(dir = compareComparables(kc, k, pk)) == 0)
dir = tieBreakOrder(k, pk);
//上面一系列都为了确定往左还是往右子树走
TreeNode<K,V> xp = p;
if ((p = (dir <= 0) ? p.left : p.right) == null) {//找到了插入的位置
x.parent = xp;//设置插入的节点父节点
if (dir <= 0)//说明插在左边
xp.left = x;
else//否则插在右边
xp.right = x;
//【标记3】进行插入后的红黑树调整
root = balanceInsertion(root, x);
break;
}
}
}
}
//【标记4】把根节点移动到数组一次(n-1)&hash就找到的下标下
moveRootToFront(tab, root);
}
【标记3】红黑树插入调整，方法会返回根节点（因为调整可能会改变根节点）。
static <K,V> TreeNode<K,V> balanceInsertion(TreeNode<K,V> root,
TreeNode<K,V> x) {
x.red = true;//插入的节点首先涂成红色
for (TreeNode<K,V> xp, xpp, xppl, xppr;;) {
if ((xp = x.parent) == null) {//说明当前树是空
x.red = false;//根节点必须黑色
return x;//插入的节点直接就是根节点
}
else if (!xp.red || (xpp = xp.parent) == null)
return root;//父节点不是红色，那就不会违反红黑树规则，直接返回原本的根节点
if (xp == (xppl = xpp.left)) {//父亲节点是祖父节点的左孩子
if ((xppr = xpp.right) != null && xppr.red) {//如果叔叔是红色
xppr.red = false;//把叔叔节点涂黑
xp.red = false;//把父亲节点涂黑
xpp.red = true;//祖父节点涂红
x = xpp;//祖父节点作为“插入节点”进入下一次循环
}
else {//叔叔节点是黑色
if (x == xp.right) {//插入节点是父节点的右孩子
//以父亲节点为支点左旋，然后更新插入节点为父亲节点【标记5】
root = rotateLeft(root, x = xp);
//更新父节点和祖父节点
xpp = (xp = x.parent) == null ? null : xp.parent;
}
//到这里插入节点一定是父节点的左孩子
if (xp != null) {
xp.red = false;//把父节点
if (xpp != null) {
xpp.red = true;//祖父节点涂成红色
root = rotateRight(root, xpp);//【标记5】以祖父节点为支点右旋
}
}
}
}
else {//父亲节点是祖父节点的右孩子
...只是做上面的镜像操作
}
}
}
【标记5】左旋右旋方法也是TreeNode的成员方法，会返回根节点（因为旋转过程中根节点可能改变了），左旋代码如下：
//root是旧的根节点，p是旋转支点
static <K,V> TreeNode<K,V> rotateLeft(TreeNode<K,V> root,
TreeNode<K,V> p) {
TreeNode<K,V> r, pp, rl;
if (p != null && (r = p.right) != null) {
if ((rl = p.right = r.left) != null)
rl.parent = p;//p的右节点的左孩子会变成p的右孩子
if ((pp = r.parent = p.parent) == null)//说明p原来是根节点
(root = r).red = false;//p的右孩子会变成根节点
else if (pp.left == p)
pp.left = r;//p原先是个左孩子
else
pp.right = r;//p原先是个右孩子
r.left = p;//p变成原先p的右节点的左孩子
p.parent = r;
}
return root;
}
【标记4】把根节点移动到数组一次(n-1)&hash就找到的下标下
static <K,V> void moveRootToFront(Node<K,V>[] tab, TreeNode<K,V> root) {
int n;
if (root != null && tab != null && (n = tab.length) > 0) {
int index = (n - 1) & root.hash;
TreeNode<K,V> first = (TreeNode<K,V>)tab[index];
if (root != first) {//root节点不是第一个节点
Node<K,V> rn;
tab[index] = root;
TreeNode<K,V> rp = root.prev;//前节点
if ((rn = root.next) != null)//
((TreeNode<K,V>)rn).prev = rp;
if (rp != null)
rp.next = rn;
//上面就是删除链表节点操作
//下面就是把节点插到链表头操作
if (first != null)
first.prev = root;
root.next = first;
root.prev = null;
}
//检查状态是否正确
assert checkInvariants(root);
}
}
二、插入操作
在调下面的putTreeVal方法前，是调了HashMap的putVal。在putVal中发现是树结构节点就会调putTreeVal方法。如果树中存在相同的key元素，putTreeVal会返回出去交由putVal方法考虑是否替换成新值，如果不存在则直接插入，并返回null。（该方法是TreeNode中的成员方法，下面的this指的是TreeNode实例）
final TreeNode<K,V> putTreeVal(HashMap<K,V> map, Node<K,V>[] tab,
int h, K k, V v) {
Class<?> kc = null;
//当hash值相同时，直接进行左右搜索。
//searched标志保证只搜索一次，如果搜索不到就不存在，那就找合适插入位置插入新元素
boolean searched = false;
//找到树的根节点
TreeNode<K,V> root = (parent != null) ? root() : this;
for (TreeNode<K,V> p = root;;) {
int dir, ph; K pk;
if ((ph = p.hash) > h)
dir = -1;//往左子树走
else if (ph < h)
dir = 1;//往右子树走
else if ((pk = p.key) == k || (k != null && k.equals(pk)))
return p;//这里找到了直接返回
else if ((kc == null &&
(kc = comparableClassFor(k)) == null) ||
(dir = compareComparables(kc, k, pk)) == 0) {
if (!searched) {
TreeNode<K,V> q, ch;
searched = true;//赋值为true，保证只搜索一次
if (((ch = p.left) != null &&
(q = ch.find(h, k, kc)) != null) ||
((ch = p.right) != null &&
(q = ch.find(h, k, kc)) != null))
return q;//找到了返回出去
}
dir = tieBreakOrder(k, pk);
}
TreeNode<K,V> xp = p;
if ((p = (dir <= 0) ? p.left : p.right) == null) {
//到这里就找到合适的插入位置了
Node<K,V> xpn = xp.next;
TreeNode<K,V> x = map.newTreeNode(h, k, v, xpn);
if (dir <= 0)
xp.left = x;
else
xp.right = x;
xp.next = x;
x.parent = x.prev = xp;
if (xpn != null)
((TreeNode<K,V>)xpn).prev = x;
//对树进行插入后的调整，然后移动根节点
moveRootToFront(tab, balanceInsertion(root, x));
return null;//返回空，说明插入了新元素
}
}
}
三、删除操作
删除操作前会找到待删除的节点，所以该方法调的是待删除节点实例的方法。（该方法是TreeNode中的成员方法，下面的this指的是TreeNode实例）
红黑树删除节分四种情况：
- 删除的节点没有孩子，那暂时先用自己代替，最后再删除
- 删除的节点只有左孩子，那用左孩子代替
- 删除的节点只有右孩子，那用右孩子代替
- 删除的节点左孩子、右孩子都不为空，那和后继节点交换，这样就转换成删除后继节点
HashMap中的红黑树实现其实是树结构和链表结构共存的，删除节点时会判断是否由于节点太少而转回链表结构，如果不转就按红黑树删除节点的方式删除节点。
一般删除红黑树节点，遇到上述的第四种情况时，是会把后继节点和删除节点的内容进行交换，然后转成前面三种情况。而HashMap的红黑树做的是把这两个节点颜色进行交换，然后位置也进行交换。
final void removeTreeNode(HashMap<K,V> map, Node<K,V>[] tab,
boolean movable) {
int n;
if (tab == null || (n = tab.length) == 0)
return;
int index = (n - 1) & hash;
TreeNode<K,V> first = (TreeNode<K,V>)tab[index], root = first, rl;
TreeNode<K,V> succ = (TreeNode<K,V>)next, pred = prev;
if (pred == null)//说明是头结点
tab[index] = first = succ;
else//从链表中断开
pred.next = succ;
if (succ != null)
succ.prev = pred;
if (first == null)//说明只有一个节点，删除后tab[index]是null了
return;
//下面就是要从树结构中删除节点
if (root.parent != null)
root = root.root();//保证root是根节点
if (root == null || root.right == null ||
(rl = root.left) == null || rl.left == null) {
tab[index] = first.untreeify(map); //【标记6】转回链表结构
return;
}
//下面是真正要从树结构中删除节点
TreeNode<K,V> p = this, pl = left, pr = right, replacement;
if (pl != null && pr != null) {//左右孩子都不是空
TreeNode<K,V> s = pr, sl;
while ((sl = s.left) != null) //找到后继节点
s = sl;
boolean c = s.red; s.red = p.red; p.red = c; // 交换颜色
TreeNode<K,V> sr = s.right;
TreeNode<K,V> pp = p.parent;
if (s == pr) { //后继节点就是待删除节点的右孩子
p.parent = s;
s.right = p;
}
else {
//下面交换位置的代码可能看的有点晕，但画一下图就懂了
TreeNode<K,V> sp = s.parent;
if ((p.parent = sp) != null) {
if (s == sp.left)
sp.left = p;
else
sp.right = p;
}
if ((s.right = pr) != null)
pr.parent = s;
}
p.left = null;
if ((p.right = sr) != null)
sr.parent = p;
if ((s.left = pl) != null)
pl.parent = s;
if ((s.parent = pp) == null)
root = s;//删除的是根节点，那后继节点直接变成根节点
else if (p == pp.left)
pp.left = s;
else
pp.right = s;
//上面交换完了，变成删除后继节点的情况了
//因为后继节点肯定没左孩子，只要判断是否有右孩子
if (sr != null)
replacement = sr;
else
replacement = p;//暂时用一下自己
}
else if (pl != null)//只有左孩子
replacement = pl;
else if (pr != null)//只有右孩子
replacement = pr;
else//左右孩子都为空
replacement = p;
if (replacement != p) {//待删除节点有孩子
TreeNode<K,V> pp = replacement.parent = p.parent;
if (pp == null)
root = replacement;
else if (p == pp.left)
pp.left = replacement;
else
pp.right = replacement;
p.left = p.right = p.parent = null;
}
【标记7】待删除节点是红色就不用调整，否则进行红黑树删除节点后的调整
TreeNode<K,V> r = p.red ? root : balanceDeletion(root, replacement);
if (replacement == p) { //待删除节点没有孩子，被暂时用了一下，下面真正移除
TreeNode<K,V> pp = p.parent;
p.parent = null;
if (pp != null) {
if (p == pp.left)
pp.left = null;
else if (p == pp.right)
pp.right = null;
}
}
if (movable)
moveRootToFront(tab, r);//移动根节点
}
【标记6】转回链表结构。由于在调用该方法前，就已经把删除的节点从链表中去掉，所以转换链表中是没有删除的节点的。
final Node<K,V> untreeify(HashMap<K,V> map) {
Node<K,V> hd = null, tl = null;
for (Node<K,V> q = this; q != null; q = q.next) {
Node<K,V> p = map.replacementNode(q, null);
if (tl == null)
hd = p;
else
tl.next = p;
tl = p;
}
return hd;
}
Node<K,V> replacementNode(Node<K,V> p, Node<K,V> next) {
return new Node<>(p.hash, p.key, p.value, next);
}
【标记7】待删除节点是红色就不用调整，否则进行红黑树删除节点后的调整
static <K,V> TreeNode<K,V> balanceDeletion(TreeNode<K,V> root,
TreeNode<K,V> x) {
//x是删除节点的替换节点
for (TreeNode<K,V> xp, xpl, xpr;;) {
if (x == null || x == root)
return root;//x是黑且又是根节点直接返回
else if ((xp = x.parent) == null) {//x是根节点
x.red = false;//涂黑
return x;//作为新根节点返回
}
else if (x.red) {//x是红色就直接涂黑返回根节点
x.red = false;
return root;
}
else if ((xpl = xp.left) == x) {//x是父亲节点的左孩子
if ((xpr = xp.right) != null && xpr.red) {//兄弟节点是红色
xpr.red = false;//把兄弟涂黑
xp.red = true;//父亲涂黑
root = rotateLeft(root, xp);//以父亲作为支点左旋
xpr = (xp = x.parent) == null ? null : xp.right;//更新兄弟节点
}
//下面开始的兄弟肯定是黑
if (xpr == null)//这算是兄弟节点为黑，且兄弟节点的两个孩子为黑的情况
x = xp;//把“替换节点”变成父亲节点继续循环
else {
//下面就要考虑兄弟节点的两个孩子都是黑，只有左孩子为红，或右孩子为红的情况
TreeNode<K,V> sl = xpr.left, sr = xpr.right;
if ((sr == null || !sr.red) &&
(sl == null || !sl.red)) {这是兄弟节点为黑，且兄弟节点的两个孩子为黑的情况
xpr.red = true;
x = xp;
}
else {
if (sr == null || !sr.red) {//兄弟节点只有左孩子为红
if (sl != null)
sl.red = false;
xpr.red = true;//兄弟节点涂红
root = rotateRight(root, xpr);//以兄弟节点为支点右旋
xpr = (xp = x.parent) == null ?
null : xp.right;//更新兄弟节点
}
//下面开始，兄弟节点的右孩子只要不是空，就肯定是红
if (xpr != null) {
xpr.red = (xp == null) ? false : xp.red;//把父节点的颜色赋给兄弟节点
if ((sr = xpr.right) != null)
sr.red = false;//兄弟节点的右孩子涂黑
}
if (xp != null) {
xp.red = false;//父亲涂黑
root = rotateLeft(root, xp);//以父亲为支点左旋
}
x = root;//赋值为根节点
}
}
}
else {
...做上面的镜像操作
}
}
}