TreeMap的源码学习

内容来自-码出高效

public class RedBlackSourceCode {

    public V put(K key, V value) {
        // t表示当前节点，记住这个很重要。先把TreeMap的根节点root引用赋值给当前节点
        TreeMap.Entry<K,V> t = root;
        // 如果当前节点为null，即是空树，新增的KV形式的节点就是根节点
        if (t == null) {
            // 看似多此一举，实际上预检了key是否可以比较
            compare(key, key); // type (and possibly null) check
            // 使用KV构造出新的Entry对象，其中第三个参数就是parent，根节点没有父节点
            root = new TreeMap.Entry<>(key, value, null);
            size = 1;
            modCount++;
            return null;
        }
        // 接收比较结果
        int cmp;
        TreeMap.Entry<K,V> parent;
        // 构造方法中置入的外部比较器
        // split comparator and comparable paths
        Comparator<? super K> cpr = comparator;
        // todo 重点步骤，根据二叉查找树的特性，找到新节点插入的合适位置
        if (cpr != null) {
            // 循环的目标，根据参数key与当前节点的key不断的进行比较
            do {
                // 当前节点赋值给父节点，故从根节点开始遍历比较
                parent = t;
                // 比较输入的参数key和当前节点key的大小
                cmp = cpr.compare(key, t.key);
                // 参数的key更小，向左边走，把当前节点引用移动到它的左子节点上
                if (cmp < 0)
                    t = t.left;
                else if (cmp > 0)
                    t = t.right;
                // 如果相等，覆盖当前节点的值，返回更新前的值。
                else
                    return t.setValue(value);
                // 如果没有相等的key，一直会遍历到NIL节点为止
            } while (t != null);
        }
        // 在没有指定比较器的情况下，调用自然排序的Comparable比较
        else {
            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);
        }
        // 创建Entry对象，并把parent置入参数
        TreeMap.Entry<K,V> e = new TreeMap.Entry<>(key, value, parent);
        // 新节点找到自己的位置，原本以为就可以安顿下来---------
        if (cmp < 0)
            // 如果比较结果小于0，则成为parent的左孩子
            parent.left = e;
        else
            parent.right = e;
        // 还需要对这个新节点进行重新着色和旋转操作，以达到平衡
        fixAfterInsertion(e);
        size++;
        modCount++;
        // 成功插入新节点后，返回null
        return null;
    }
    /** From CLR */
    private void fixAfterInsertion(TreeMap.Entry<K,V> x) {
        // 虽然内部类Entry的属性color默认为黑色，但是新节点一律先赋值为红色
        x.color = RED;

        /*新节点是根节点或者其父节点（简称父亲）为黑色
        *插入红色节点并不会破坏红黑树的性质，无须调整
        * x值改变的过程是在不断地向上游遍历,直到父亲为黑色，或者到达根节点
        * */
        while (x != null && x != root && x.parent.color == RED) {
            // 如果父亲是其父节点（爷爷）的左子节点
            if (parentOf(x) == leftOf(parentOf(parentOf(x)))) {
                // 这时，得到爷爷的右子节点（右叔）的脸色
                TreeMap.Entry<K,V> y = rightOf(parentOf(parentOf(x)));
                // 如果右叔是红色的，此时通过局部颜色调整，就可以使子树继续满足红黑树性质
                // todo 第一处
                if (colorOf(y) == RED) {
                    // 父亲设置为黑色
                    setColor(parentOf(x), BLACK);
                    // 右叔设置为黑色
                    setColor(y, BLACK);
                    // 爷爷设置为红色
                    setColor(parentOf(parentOf(x)), RED);
                    // 爷爷成为新的节点，进入到下一轮循环
                    x = parentOf(parentOf(x));
                    // 如果右叔是黑色，则需要加入旋转
                } else {
                    // 对父亲做一次左旋转操作，红色的父亲会沉入其左侧位置
                    // 将父亲赋值给x
                    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 {
                TreeMap.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;
    }

    // 左旋转
    private void rotateLeft(TreeMap.Entry<K,V> p) {
        // 如果参数节点不是NIL节点
        if (p != null) {
            // 获取p的右子节点r
            TreeMap.Entry<K,V> r = p.right;
            // 将 r的左子树 赋值给 p的右子树
            p.right = r.left;
            //
            if (r.left != null)
                r.left.parent = p;
            r.parent = p.parent;
            if (p.parent == null)
                root = r;
            else if (p.parent.left == p)
                p.parent.left = r;
            else
                p.parent.right = r;
            // 将p设置为r的左子树，将r设置为p的父亲
            r.left = p;
            p.parent = r;
        }
    }
}