从TreeMap看红黑树实现

199 阅读5分钟
代码版本 jdk 1.8

新增

put 方法

通过TreeMap的put方法入手

public V put(K key, V value) {
    Entry<K,V> t = root;
    //如果root节点为空,证明树还没有初始化
    if (t == null) {
        //比较key的大小
        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
    Comparator<? super K> 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);
    }
    //如果没有初始化Comparator
    else {
        //通过这个地方可以发现,没有初始化Comparator 将不允许传入控制
        if (key == null)
            throw new NullPointerException();
        @SuppressWarnings("unchecked")
            Comparable<? super K> k = (Comparable<? super K>) key;
        //do while 作用 : 从树的根节点一节节往下寻找,
        //如果值大于当前寻找的节点,继续寻找当前节点的右节点,
        //如果小于,继续寻找树的右节点,直到寻找到子节点为空的节点
        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);
    }
    //创建一个新的节点,并判断此节点该放在partner节点的左节点还是右节点
    Entry<K,V> e = new Entry<>(key, value, parent);
    if (cmp < 0)
        parent.left = e;
    else
        parent.right = e;
    //如果是普通的二叉树,此阶段就结束了,但是TreeMap是红黑树,为了满足红黑树的特性,
    //执行此方法,做树的修复
    fixAfterInsertion(e);
    size++;
    modCount++;
    return null;
}

此处代码的功能就是将节点塞入到树指定的地方,接下来看fixAfterInsertion(e)方法,看怎么修复的

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;
}

我们看一下红黑色的特性,结合上面的代码看一下Java是怎么具体实现它的

  1. 节点是红色或黑色。
  2. 根是黑色。
  3. 所有叶子都是黑色(叶子是NIL节点)。
  4. 每个红色节点必须有两个黑色的子节点。(从每个叶子到根的所有路径上不能有两个连续的红色节点。)
  5. 从任一节点到其每个叶子的所有简单路径都包含相同数目的黑色节点。

1、2、3 每个Entry创建的时候,就会有默认颜色BLACK,之后的大部分篇幅都是为了实现4、5点 fixAfterInsertion(e),分为一下集中情况

  • 父节点不为红色 所有条件都不会被破坏,不继续做操作
  • 父节点为红色
    • 当前节点的父节点是祖父节点左节点
      • 叔叔节点为红色 case1
      • 叔叔节点不为红色 case2
    • 当前节点的父节点是祖父节点的右节点
      • 叔叔节点为红色 case3
      • 叔叔节点不为红色 case4

总结一下各种case的处理情况 图片中的所有最下级叶子节点nul省略

case1,case3可以合并,他们做的操作都一样,只不过是找叔叔节点的方式不一样

设置当前节点的父节点,叔叔节点为黑色,设置祖父节点为红色,将当前节点设置为祖父节点

上图展示了变化过程,此时可能节点5的父节点也为红色,需要将当前节点置为节点5 递归处理这种情况

case2 当前节点的父节点为祖父节点的左节点,并且叔叔节点不为红色

2.1 如果当前节点为父节点的右节点,将当前节点置为父节点,然后左旋当前节点,
 如果为左节点,直接进行2.2流程的处理

这样经过左旋之后就变成了下面这种情况了,之后会以1为当前节点操作

2.2 将当前节点父节点设置为黑色,设置祖父节点为红色,右旋当前节点的祖父节点

case4 当前节点的父节点为祖父节点的右节点,并且叔叔节点不为红色

4.1 如果当前节点为父节点的左节点,将当前节点置为父节点,然后右旋当前节点(右旋父节点),如果为右节点,直接进行4.2操作

经过上面处理,然后将8节点作为当前节点,继续进行4.2操作

4.2 将当前节点父节点设置为黑色,祖父节点设置为红色,左旋当前节点的祖父节点

删除 待补充