TreeMap学习

树节点——TreeMap的节点结构

static final class Entry<K,V> implements Map.Entry<K,V> {
    K key;
    V value;
    Entry<K,V> left;
    Entry<K,V> right;
    Entry<K,V> parent;
    boolean color = BLACK;
    
    Entry(K key, V value, Entry<K,V> parent) {
        this.key = key;
        this.value = value;
        this.parent = parent;
    }
}

一个TreeMap节点主要包括键（key）、值（value）、左子节点（left）、右子节点（right）、父节点（parent）、节点颜色（color，默认是黑色）六个属性，构造方法需要键、值和父节点。

后面主要从三个方法具体解析TreeMap结构，put、get、remove。

put方法——如何构建一个TreeMap

如何从无到有的构建一个TreeMap？

常规使用TreeMap以及其他map的过程中，我们通常会用无参构造方法new一个map，然后将数据逐个put的方式，我们主要分析这种情况。

1、TreeMap类的四个属性

private final Comparator<? super K> comparator;  //TreeMap的比较器
private transient Entry<K,V> root;  //根节点
private transient int size = 0; //TreeMap下的节点数
private transient int modCount = 0; //TreeMap下的操作数

2、TreeMap的无参构造方法

public TreeMap() {
    comparator = null;
}

TreeMap的比较器在调用无参构造时初始化为null。

3、TreeMap下的put方法

public V put(K key, V value) {
    Entry<K,V> t = root;   //声明节点t指向root节点，不直接操作root节点
    if (t == null) {       //如果根节点为空，则进行初始化
        compare(key, key); //根据是否有比较器进行比较
        root = new Entry<>(key, value, null);  //插入的节点作为root节点
        size = 1;
        modCount++;
        return null;
    }
    int cmp;           //根节点不为空，cmp为比较器的返回值，-1，0,1三个值
    Entry<K,V> parent; //记录查找过程中的父节点，用于初始化节点
    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); //查找到key相同的键，直接将value存入并返回旧值
        } while (t != null);
    }
    else { //如果比较器为空则根据key的比较方式进行比较
        if (key == null) //TreeMap在未自定义比较器的时候是不支持key为空指针的
            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<K,V> e = new Entry<>(key, value, parent);
    if (cmp < 0) //根据cmp的值决定插入到父节点的左边还是右边
        parent.left = e;
    else
        parent.right = e;
    fixAfterInsertion(e);  //插入后的调整,红黑树最核心的方法
    size++;
    modCount++;
    return null;
}

通过上述分析，TreeMap的插入操作，主要是通过比较器判断是更新值还是新增节点，对于新增节点进行插入并作出平衡调整，下面重点分析fixAfterInsertion方法

private void fixAfterInsertion(Entry<K,V> x) {
    //首先保证插入节点为红色。此处为红色的原因是，因为红黑树的一条特性，从根节点到任何一个叶子节点的路径上的黑节点数相同，如果插入的是红色节点，则不会破坏此规则。
    //回顾其他红黑树特性，二叉搜索树和节点带颜色两个特点都满足，那么只需要保证，红节点的子节点一定不是红色即可
    //所以只有一种情况需要调整，就是x节点的父节点为红色的情况，此时祖父节点一定为黑或者空
    x.color = RED;
    //x如果为空或者为根节点则直接不用循环
    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) {//情况1：父节点和叔节点均为红
                //此时同时将父节点和叔节点变黑，不会影响特性3
                setColor(parentOf(x), BLACK);
                setColor(y, BLACK);
                //由于整体新增了一个黑节点，只能将祖父节点涂红来保证平衡
                setColor(parentOf(parentOf(x)), RED);
                //为了避免祖父节点和其父节点出现连续红节点的情况，将x指向祖父节点进行迭代判断
                x = parentOf(parentOf(x));
            } else { //父节点为红，叔节点为黑
                if (x == rightOf(parentOf(x))) {//情况2：x为父节点的右子节点
                    x = parentOf(x);
                    rotateLeft(x);
                }
                //情况3：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;
}

get方法——如何获取TreeMap中指定的节点值

public V get(Object key) {
    Entry<K,V> p = getEntry(key);
    return (p==null ? null : p.value);
}

调用getEntry方法，此处的代码很熟悉，和put方法中的查找方法一模一样。

final Entry<K,V> getEntry(Object key) {
    // Offload comparator-based version for sake of performance
    if (comparator != null)
        return getEntryUsingComparator(key);
    if (key == null)
        throw new NullPointerException();
    @SuppressWarnings("unchecked")
        Comparable<? super K> k = (Comparable<? super K>) key;
    Entry<K,V> p = root;
    while (p != null) {
        int cmp = k.compareTo(p.key);
        if (cmp < 0)
            p = p.left;
        else if (cmp > 0)
            p = p.right;
        else
            return p;
    }
    return null;
}

remove方法——如何获取并删除节点中的值

remove方法和get方法一样，首先找到指定的值，然后进行删除，重点是删除操作以及删除后的平衡调整

public V remove(Object key) {
    Entry<K,V> p = getEntry(key);
    if (p == null)
        return null;

    V oldValue = p.value;
    deleteEntry(p);
    return oldValue;
}

重点分析deleteEntry方法

private void deleteEntry(Entry<K,V> p) {
    //操作次数和大小处理
    modCount++;
    size--;
    
    
    if (p.left != null && p.right != null) { //如果p有左右子节点，后继节点一定在其右子树上
        //查找到后继节点，并将后继节点的值传给当前节点，然后p指向后继节点，这样删掉后继节点即可
        Entry<K,V> s = successor(p);
        p.key = s.key;
        p.value = s.value;
        p = s;
    } 

    //寻找到replacement节点，如果左节点不为空，则为左节点，左节点为空，则为右节点
    Entry<K,V> replacement = (p.left != null ? p.left : p.right);
    
    //逻辑走到这个地方可以简单的先分析一下：
    //p如果有两个子节点，则寻找后继节点，后继节点的左子节点肯定为空，所以p变成一个最多只有一个子节点的节点
    //所以这个替换节点replacement就是单子节点中的一个，或者为一个空节点
    
    //不为空的时候，也就是p节点的单子节点，其实此时就是一个单链表的删除
    if (replacement != 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;
        p.left = p.right = p.parent = null;
        //删除了p节点，如果p为红，则不会影响红黑树的特性，所以，如果为黑的时候才会进行调整
        if (p.color == BLACK)
            fixAfterDeletion(replacement); 
    } else if (p.parent == null) { //逻辑到达这个地方，说明p没有子节点且为根直接，直接将root置空即可
        root = null;
    } else { //如果不为根节点，则直接进行删除
        if (p.color == BLACK)//如果节点颜色为黑，需要进行调整
            fixAfterDeletion(p);
        if (p.parent != null) { //直接将p的父节点的子节点指向变为空，相当于删除链表的最后一个节点
            if (p == p.parent.left)
                p.parent.left = null;
            else if (p == p.parent.right)
                p.parent.right = null;
            p.parent = null;
        }
    }
}

经过上面的分析，重点放在删除后调整这个方法上fixAfterDeletion，看这个方法之前，主要是两个情况会调用这个方法

1、p节点只有一个子节点，删除p节点后，以p节点的下一个节点为入参调用

2、p节点无子节点，以p节点为入参调用方法后，删除p节点

private void fixAfterDeletion(Entry<K,V> x) {
    //出循环的条件：x指向了根节点|x的颜色为红色
    //根据上述调用删除调整的两个情况分析
    //第一种情况：p节点只有一个子节点，且删除的节点为黑色才会调用此方法。如果替代节点为红色，直接将替代节点变黑即可;如果替代节点为黑色，则替代节点所属的子树少了一个黑色的节点
    //第二种情况，p无子节点且为黑色，此时的p没有被删除，但后续删除后对应子树会减少一个黑色节点
    //所以调整方案就是增加当前子树一个黑色节点
    while (x != root && colorOf(x) == BLACK) {...}
        if (x == leftOf(parentOf(x))) { //x为父节点的左节点
            Entry<K,V> sib = rightOf(parentOf(x)); //兄弟节点
            if (colorOf(sib) == RED) { //情况1：兄弟节点为红色，推出父节点和sib的子节点都为黑色
                //此调整之前说过，不会改变红黑树的结构
                setColor(sib, BLACK);
                setColor(parentOf(x), RED);
                rotateLeft(parentOf(x));
                sib = rightOf(parentOf(x));
            }
            //情况②上面的情况的最终结果也是变成这个情况
            //直接将兄弟节点变为红色，x指向x的父节点，如果是从上面情况到达的，x的parent节点一定是红色，首先是将兄弟节点涂红减少兄弟子树的一个黑色节点，然后跳出循环，将父节点涂黑，这样给整个删除子树增加了一个黑色节点，从而完成调整
            //如果直接是情况2，则直接将右兄弟节点涂红，然后x指向x的父节点，这样如果x的父节点是红色，则直接涂黑退出即可，但是如果是黑色节点，则相当于是父节点和其兄弟节点进行对比
            //情况234：兄弟节点均为黑色
            if (colorOf(leftOf(sib))  == BLACK &&
                colorOf(rightOf(sib)) == BLACK) {//情况2：兄弟节点的子节点均为黑色
                setColor(sib, RED);
                x = parentOf(x);
            } else {
                if (colorOf(rightOf(sib)) == BLACK) {//情况3：兄弟节点的左子节点为红色
                    setColor(leftOf(sib), BLACK);
                    setColor(sib, RED);
                    rotateRight(sib);
                    sib = rightOf(parentOf(x));
                }
                setColor(sib, colorOf(parentOf(x)));//情况4：兄弟节点的右子节点为红色
                setColor(parentOf(x), BLACK);
                setColor(rightOf(sib), BLACK);
                rotateLeft(parentOf(x));
                x = root;
                //此处的调整，和之前的调整不太一样，之前的调整都是保证调整前后的左右子树的黑节点个数不一样，此处的核心就是对父节点的左旋，第一个颜色改变，是保证左旋前后，节点的颜色不会发生变化，第二个颜色改变是保证左旋后的左节点能增加一个黑节点，右节点由于左旋会减少一个黑节点，第三个颜色改变，是保证左旋后的右子节点黑色节点数增加一个，这样就完成了整个的调整，保证删除子树增加了一个黑色节点，即可以直接退出循环
            }
        } else { // symmetric
            Entry<K,V> 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);
}