萌新从TreeMap学习红黑树

567 阅读13分钟

萌新学习数据结构挺久的了,常用数据结构都可以手撕,而平衡树只是了解原理,撕不出来,看各种博客文章也看得晕头转向的。
之前看《算法》红皮书学习了左偏红黑树,这次从JDK的TreeMap来分析下常规红黑树。 阅读需要有二叉查找树的知识背景

1.红黑树的基本性质

出自《算法导论》

  1. 每个节点是红色或者黑色的
  2. 根节点黑
  3. 叶节点Nil黑
  4. 父子不能都是红色的
  5. 从一个节点到其子孙节点的所有路径上,黑色节点的数目是相同的

2.从TreeMap中提取红黑树相关的代码

因为TreeMap中有很多集合相关的操作,原代码长度上千行,看得眼花了。所以这里把其中和红黑树相关的部分提取出来分析。

基本类及其属性

节点类

这里为了方便,把其中的泛型部分简化成int

    private static class Node implements Comparable<Node> {
        int key;
        int val;
        boolean color = BLACK;
        Node left, right, parent;

        Node(int key, int val, Node parent) {
            this.key = key;
            this.val = val;
            this.parent = parent;
        }

        @Override
        public int compareTo(Node o) {
            return this.key - o.key;
        }
    }

红黑树类

//红黑树
public class RedBlackTree{
    //常量定义
    private static final boolean RED = false;
    private static final boolean BLACK = true;
    
    private Node root;
    private int size;
    // 之前的节点类
    private static class Node implements Comparable<Node> {...}
}

工具方法

颜色相关操作

    private static boolean colorOf(Node p) {
        return (p == null ? BLACK : p.color);
    }

    private static void setColor(Node p, boolean c) {
        if (p != null)
            p.color = c;
    }

节点关系相关操作

    private static Node parentOf(Node p) {
        return (p == null ? null : p.parent);
    }

    private static Node leftOf(Node p) {
        return (p == null) ? null : p.left;
    }

    private static Node rightOf(Node p) {
        return (p == null) ? null : p.right;
    }
    
    private Node successor(Node tmp) {
    //后继节点的查找
        if (tmp == null) {
            return null;
        }
        if (tmp.right != null) {
            Node p = tmp.right;
            while (p.left != null) {
                p = p.left;
            }
            return p;
        } else {
            Node p = tmp.parent;
            Node ch = tmp;
            while (p != null && ch == p.right) {
                ch = p;
                p = p.parent;
            }
            return p;
        }
    }

旋转相关操作

所谓左旋,就是对原节点N,让N的右子节点R代替N的位置,
N成为R的左子节点,至于他们的其他子节点,看着办就好。
从观感上有N下位成左子节点,N的右子节点上位之感,故谓之左旋。
右旋反之。
非常简单的操作,无须多言。

    private void rotateLeft(Node tmp) {
       if (tmp == null) {
           return;
       }

       Node r = tmp.right;
       if (r == null) {
           return;
       }

       tmp.right = r.left;
       if (r.left != null) {
           r.left.parent = tmp;
       }

       r.parent = tmp.parent;
       if (tmp.parent == null) {
           root = r;
       } else if (tmp.parent.left == tmp) {
           tmp.parent.left = r;
       } else {
           tmp.parent.right = r;
       }

       r.left = tmp;
       tmp.parent = r;
   }

   private void rotateRight(Node tmp) {
       if (tmp == null) {
           return;
       }
       Node l = tmp.left;
       if (l == null) {
           return;
       }
       tmp.left = l.right;
       if (l.right != null) {
           l.parent = tmp;
       }

       l.parent = tmp.parent;

       if (tmp.parent == null) {
           root = l;
       } else if (tmp == tmp.parent.left) {
           tmp.parent.left = l;
       } else {
           tmp.parent.right = l;
       }

       l.right = tmp;
       tmp.parent = l;
   }

插入

整体方法

  • 空树直接插入到root
  • 非空树插入后,看是否建立了新的节点,有新的节点就需要修复RBT的性质
  • 这里有两个方法需要进一步分析,insert 和 fixAfterInsertion
    public void put(int key, int val) {
        if (this.root == null) {
            //对于空树,我们直接插入就好了,满足RBT的各种性质
            this.root = new Node(key, val, null);
            size = 1;
        } else {
            Node newNode = insert(root, key, val);
            //newNode为空,就是key存在,这时候没有插入新节点
            //非空的话,插入新节点,需要考虑修复被破坏的RBT性质
            if (newNode != null) {
                  fixAfterInsertion(newNode);
            }
        }
    }

insert 方法分析

其实与常规二叉树的操作一样的

    private Node insert(Node root, int key, int val) {
        Node tmp = root, parent = tmp;
        while (tmp != null) {
            parent = tmp;
            int cmp = tmp.key - key;
            if (cmp < 0) {
                tmp = tmp.left;
            } else if (cmp > 0) {
                tmp = tmp.right;
            } else {
                // key已经存在,直接修改val后返回就好了
                tmp.val = val;
                return null;
            }
        }
        
        //tmp为null, parent是上一轮的tmp
        //只需要把节点插入到这个parent下面
        int cmp = parent.key - key;
        Node newNode = new Node(key, val, parent);
        if (cmp > 0) {
            parent.right = newNode;
        } else {
            parent.left = newNode;
        }
        size++;
        return newNode;
    }

fixAfterInsertion分析,重头戏

方法名分析:插入之后的修复。修复必然是因为某些性质被破坏,这里需要通过一些操作来还RBT的性质。

几个关键点

  1. 首先新插入的节点设为红色,这样可以保证性质5,我们只需要维护性质4(父子节点不能都是红节点)就好了
  2. 循环条件为当前节点的父节点为红色,因为只有这种情况下我们才需要维护被破坏的性质4,否则的话符合RBT的性质,可以返回了
  3. 循环中的处理分成两种,按当前节点Tmp的父节点Parent是Tmp的祖父节点的左子节点还是右子节点来分,其内部的处理方式是镜像的,我们只需要分析其中一种就好(第一个if)
  4. 先获取当前节点的叔节点,其父亲的兄弟。 兄弟节点可能的几种状态:
    • 红色
    • 黑色
    • 不存在
  5. 当叔节点为RED时,即代码中第二个if处,父节点也为RED,祖节点必为BLACK。
    我们把叔节点和父节点都变为BLACK,祖节点变为RED。
    这样操作后:
    一方面,对于父节点和叔节点及他们的子节点来说,从根节点到当前节点路径上的黑色节点数目没有变化。不会破坏RBT的性质。
    另一方面,父节点变成黑色,当前节点和父节点不再都是红色了,性质四被修复。
    这个修复只是针对当前节点来说的,因为祖节点变成了红色,祖节点的性质4可能被破坏,所以需要把tmp指向祖节点,继续下一轮循环。

变成

这个操作,可以参考算法红皮书中的flip方法

  1. 当叔节点为黑色或不存在时,第二个else那里
    6.1 判断当前节点是左子节点,还是右子节点,如果是右子节点,那么使用旋转操作把他变成左子节点,方便后面的统一处理

6.2 正式处理:

很好,通过旋转,成功把一个红色节点甩锅给叔节点了,现在性质4又得到了满足。
注意在旋转前,祖节点到tmp的路径上只有一个黑色节点,祖节点到右边的孙子的路径上有两个黑色节点,而旋转后黑色节点的数目保持不变,从而维护了RBT的性质

    private void fixAfterInsertion(Node insertNode) {
        setColor(insertNode, RED);
        Node tmp = insertNode;
        while (tmp != null && tmp != root && tmp.parent.color == RED) {
             //第一个if
            if (parentOf(tmp) == leftOf(parentOf(parentOf(tmp)))) {
                Node uncle = rightOf(parentOf(parentOf(tmp)));//tmp的叔节点
                if (colorOf(uncle) == RED) {
                    //第二个if
                    //可参考红皮书的翻转操作
                    setColor(parentOf(tmp), BLACK);
                    setColor(uncle, BLACK);
                    setColor(parentOf(uncle), RED);
                    tmp = parentOf(uncle);
                } else {//第二个else,叔节点为黑色或不存在
                    if (tmp == rightOf(parentOf(tmp))) {//第三个if
                        //若tmp为右节点,那么通过旋转操作,使tmp指向左子节点,方便下面的统一操作
                        tmp = parentOf(tmp);
                        rotateLeft(tmp);
                    }
                    setColor(parentOf(tmp), BLACK);
                    setColor(parentOf(parentOf(tmp)), RED);
                    rotateRight(parentOf(parentOf(tmp)));
                }
            }
            //else中的内容为第一个if的镜像
            else {
                Node uncle = leftOf(parentOf(parentOf(tmp)));//tmp的左叔节点
                if (colorOf(uncle) == RED) {
                    setColor(parentOf(tmp), BLACK);
                    setColor(uncle, BLACK);
                    setColor(parentOf(uncle), RED);
                    tmp = parentOf(uncle);
                } else {
                    if (tmp == leftOf(parentOf(tmp))) {
                        tmp = parentOf(tmp);
                        rotateRight(tmp);
                    }
                    setColor(parentOf(tmp), BLACK);
                    setColor(parentOf(parentOf(tmp)), RED);
                    rotateLeft(parentOf(parentOf(tmp)));
                }
            }
        }
        setColor(root, BLACK);
    }

从编程的角度理解,这里循环的作用,是把“新加了一个红色节点”这一事件逐层向上传递
在传递过程中,可能某一层可以处理这一事件,那么他就处理,然后终止这一事件的传递。如果处理不了这一事件,他就通过一系列转换,把这个事件转换成上层要处理的问题。
这样递归到root就好了。

查找

其实与常规二叉树查找一样

    public Node get(int key) {
        Node tmp = root;
        while (tmp != null) {
            int cmp = tmp.key - key;
            if (cmp > 0) {
                tmp = tmp.right;
            } else if (cmp < 0) {
                tmp = tmp.left;
            } else {
                return tmp;
            }
        }
        return null;
    }

删除

基本方法

    public void delete(int key) {
        Node node = get(key);
        if (node != null) {
            size--;
            deleteNode(node);
        }
    }

deleteNode

关键点分析:

  1. 删除红色节点不会破坏RBT的性质,但是删除黑色节点会破坏性质5,所以删除黑色节点后需要调用fixAfterDeletion方法来修复性质5,具体后面分析

  2. 先是各种非空判断。

  3. 语句1:对于要删除有两个子节点的节点Tmp,我们先找的其后继节点s,然后把s提到Tmp的位置,转为删除s就好了。
    因为Tmp有两个子节点,可知s必定存在于Tmp的右子树中,并且s没有子节点或者只有一个子节点。
    这样我们就把所有的删除操作都归纳到删除叶子节点或者只有一个子节点两种情况下了。

  4. 对于只有一个子节点的情况,我们在语句2中处理。 主要步骤就是找到这个子节点,然后子节点登基大宝,原节点被忘却。
    但是原节点若为黑色,那么这条路径下所有节点的路径都少了一个黑色节点,不符合性质5了,所以要进行修复。

  5. 对于删除叶子节点的情况,我们在语句3中处理。 该节点为红色,我们就直接断开各种连接就好了; 若该节点为黑色,我们还需要先进行fixAfterDeletion。

    private void deleteNode(Node node) {
        if (node == null) {
            return;
        }
        if (node.parent == null) {
            root = null;
            return;
        }
        
        //语句1
        if (node.left != null && node.right != null) {
            Node s = successor(node);
            node.key = s.key;
            node.val = s.val;
            node = s;
        }
        
        //语句2
        if (node.left != null || node.right != null) {
            //找到这个唯一的子节点
            Node replacement = node.left == null ? node.right : node.left;
            
            //把这个子节点顶上去,原节点的各种连接都被这个子节点所占据
            replacement.parent = node.parent;
            if (node.parent == null) {
                root = replacement;
            } else if (node == node.parent.left) {
                node.parent.left = replacement;
            } else {
                node.parent.right = replacement;
            }
            node.left = null;
            node.right = null;
            node.parent = null; 
            
            //老节点若为黑色,需要修复
            if (node.color == BLACK) {
                fixAfterDeletion(replacement);
            }

        } else {//语句3
              if (node.parent == null) {
                root = null;
            } else {
                if (node.color == BLACK) {
                    fixAfterDeletion(node);
                 }
                //断开连接
                if (node.parent != null) {
                    if (node == node.parent.left) {
                        node.parent.left = null;
                    } else if (node == node.parent.right) {
                        node.parent.right = null;
                    }
                    node.parent = null;
                }
            }
        }
    }

fixAfterDeletion

如果删除体现了“新王上位老王败溃”的无情,那么修复则体现了“兄弟就是拿来坑的”的暖暖亲情。
删除之后,若不满足RBT的性质,只会是不满足性质5。即和其兄弟比起来,在路径上少了一个黑色节点。
所以这个修复的关键是想办法从兄弟那里找一个红色节点变成黑色,通过parent传递过来,从而达到平衡。实在搞不到了,那么就把兄弟也减少一个黑色节点(黑变红),然后把问题交给parent去处理,充分体现了高超的甩锅水平。
关键点分析

  1. 循环的条件就是当前节点是黑色,因为只有当前节点是黑色的时候,才会对路径上缺少黑色节点这一问题束手无策;若当前节点是红色,直接把他变成黑色,问题就解决了。这也是循环后要进行的一个处理。
  2. 和插入后的修复类似,这里也是镜像条件判断,分成两种情况来分析,即当前节点是父节点的左子节点还是右子节点,我们只需要分析当前节点是左子节点这一种情况就好了。
  3. 我们首先获取兄弟节点bro,必定存在,不然不满足性质5。对于bro,我们分成bro颜色为红色和黑色来处理。
  4. 首先如果bro的颜色是红色,我们通过一系列操作把他变成黑色的,见代码1。
    即在不改变现有的路径黑色节点数的前提下,通过旋转和换色操作,把兄弟节点变成黑的,从而方便之后的统一处理。
  5. 无法从兄弟处找到红色节点:对于代码2,此处bro的颜色必定为黑色,若bro的两个子节点颜色也都是黑色,我们就没办法骗过来一个红色节点,所以这时候只能把问题交给上一级去解决。

先把兄弟变成红色,让bro路径下的黑色节点数目和tmp路径下的黑色节点数目一样。现在问题变成了parent路径下的黑色节点数目比parent的兄弟路径下的黑色节点数目少一个。
6. 从bro处借来一个节点:从时bro必定为黑色,且其子节点中至少存在一个红色,我们把这个红色变黑,bro就多了一个黑色,再把多的这个通过parent传递过来,来修复tmp路径下的缺憾。 代码3,4
第一步: 为了统一处理,我们只处理bro的右子节点是红色的情况,但如果他的右子节点是黑色怎么办?转化!代码3

注意这里的parent颜色不一定为黑,黑红都可能
第二步:借节点,代码4
把兄弟提上去,兄弟的右子节点变黑(原先为红,第一步中保证为红) 把原先的parent变黑拽下来,给tmp的黑节点充数

private void fixAfterDeletion(Node tmp) {
        while (tmp != root && colorOf(tmp) == BLACK) {
            if (tmp == leftOf(parentOf(tmp))) {
                //获取兄弟节点
                Node bro = rightOf(parentOf(tmp));
                
                //代码1,把兄弟变成黑色
                if (colorOf(bro) == RED) {
                    //parent一定为黑
                    setColor(bro, BLACK);
                    setColor(parentOf(tmp), RED);
                    rotateLeft(parentOf(tmp));
                    bro = rightOf(parentOf(tmp));
                }

                //代码2,bro的颜色必定为黑色
                if (colorOf(leftOf(bro)) == BLACK && colorOf(rightOf(bro)) == BLACK) {
                    setColor(bro, RED); 
                    tmp = parentOf(tmp);//向上一层传递事件
                } 
                else {
                //代码3,bro为黑色,并且bro至少有一个红子节点
                    if (colorOf(rightOf(bro)) == BLACK) {
                        setColor(leftOf(bro), BLACK);//这种情况想bro的左子节点必定为RED
                        setColor(bro, RED);
                        rotateRight(bro);
                        bro = rightOf(parentOf(tmp));
                        //成功把bro的right转成了红色
                    }
                    
                    //代码4这一步的作用是从bro那里借来一个黑色节点
                    setColor(bro, colorOf(parentOf(tmp)));
                    setColor(parentOf(tmp), BLACK);
                    setColor(rightOf(bro), BLACK);
                    rotateLeft(parentOf(tmp));
                    
                    tmp = root;//退出
                }
            } else {//镜像操作,没啥好说的
                Node bro = leftOf(parentOf(tmp));

                if (colorOf(bro) == RED) {
                    setColor(bro, BLACK);
                    setColor(parentOf(tmp), RED);
                    rotateRight(parentOf(tmp));
                    bro = leftOf(parentOf(tmp));
                }
                if (colorOf(rightOf(bro)) == BLACK && colorOf(leftOf(bro)) == BLACK) {
                    setColor(bro, RED);
                    tmp = parentOf(tmp);
                } else {
                    if (colorOf(leftOf(bro)) == BLACK) {
                        setColor(rightOf(bro), BLACK);
                        setColor(bro, RED);
                        rotateLeft(bro);
                        bro = leftOf(parentOf(tmp));
                    }
                    setColor(bro, colorOf(parentOf(tmp)));
                    setColor(parentOf(tmp), BLACK);
                    setColor(leftOf(bro), BLACK);
                    rotateRight(parentOf(tmp));
                    tmp = root;
                }
            }
        }
        setColor(tmp, BLACK);
    }

理解这一步的关键就是:通过循环,把少了一个黑色节点这一事件逐层向上传递,直到被某一层处理。具体处理方法呢?就是在bro的子节点里找到一个RED。
疑问为什么不在bro为RED的时候直接把bro变成BLACK,通过parent传递过来呢?因为bro的子节点的黑色路径数目会变,如下图

3.总结

问题的等效转换
初看可能会觉得,各种情况极其复杂,然而他们通过等效变换(换色+旋转,保持各个节点路径的黑色节点数目不变),将各种复杂情况归纳为两三种简单情况,对这两三种简单情况,又分为在本层解决,或者在本层局部解决把问题推到上一层这两种处理方式,从而保持或者修复RBT的各种性质。 插入总结

  • 父红,叔节点为红色。执行flip操作,把问题推到上一层
  • 父红,叔节点黑色,通过换色和旋转,把一个红色节点转移到叔节点那棵树上

删除总结

  • 先把兄弟节点变成黑色的
  • 兄弟节点的左右子节点都是黑色:处理不了,把兄弟也变红,保持tmp和兄弟的平衡,至于少了一个黑色的事情,交给parent去处理
  • 兄弟节点的左右子节点存在红色,把红色统一转移到右边,然后parent下降到tmp这里并变黑,补足黑节点,兄弟上升为新的parent,兄弟的右子节点变黑