Java集合之TreeMap源码解析

上期回顾

上期我从树型结构谈到了红黑树的概念以及自平衡的各种变化（指路上期←戳），本期我将会对TreeMap结合红黑树理论进行解读。

首先，我们先来回忆一下红黑树的5条基本规则。

1.结点是红色或者黑色，

2.根结点为黑色，

3.每个叶子结点都是黑色的空结点，

4.每个红色结点的两个子结点都是黑色，

5.从任何一个结点到叶子结点的所有路径的黑色结点的个数相同。

 

TreeMap的基本结构

我们先来看一下TreeMap的成员属性，

    private final Comparator<? super K> comparator;
    private transient Entry<K,V> root;
    private transient int size = 0;
    private transient int modCount = 0;

第1行，比较器，主要是用来维持树结构的有序，可以为空，如果为空，那按照key的自然规则进行排序，

第2行，Entry的引用root，是树型结构的根，

第3行，记录map中数据元素的个数，

第4行，记录map中数据更新的次数，

 

Entry是TreeMap的一个静态内部类，一个Entry对应一个树型结构中的结点，下面我们来看一下Entry的成员属性和构造函数，

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

第2行，键

第3行，值

第4行，左结点的引用

第5行，右结点的引用

第6行，父结点的引用

第7行，当前结点默认为黑色结点

第9行，构造函数。初始化键、值、父结点引用

 

红黑树中的第1条规则是结点是红色或者黑色，在TreeMap中声明了两个常量表示

private static final boolean RED   = false;
private static final boolean BLACK = true;

false表示红色，true表示黑色，根结点是黑色，因此是true。

 

TreeMap的put方法解读

下面我们解读一下map在新增元素时发生哪些动作？我们以一个public修饰的静态方法，返回值为void，参数为String数组的main方法

作为程序的切入口，如下一段代码，

public static void main(String[] args) {
        TreeMap<Integer, Integer> map = new TreeMap<Integer, Integer>();
        map.put(50,3001);
        map.put(30,3002);
        map.put(70,3003);
}

第2行，new TreeMap<Integer, Integer>()，是一个无参的构造函数，初始化比较器为null，

第3行，新增一个键为50，值为3001的元素，

我们来跟一下put方法，源代码如下，

 

public V put(K key, V value) {
        Entry<K,V> t = root;
        if (t == null) {
            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<? 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);
        }
        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<K,V> e = new Entry<>(key, value, parent);
        if (cmp < 0)
            parent.left = e;
        else
            parent.right = e;
        fixAfterInsertion(e);
        size++;
        modCount++;
        return null;
}

第2行，把变量t指向根结点的引用，

第3行，判断根结点是否为空，如果为空说明，当前数据元素时首次新增，

第4行，调用自定义的比较器或者自然比较器，自身与自身做比较，旨在判断key是否是null，如果为空，则抛出NPE，

因此，TreeMap中的Key不允许为空。

第6-9行，把root引用指向当前元素，并且树结构的数量size更新为1，modCount自增，结束，返回值为null，

 

第11行，声明一个int类型的cmp，用来记录比较器的返回值

第12行，声明一个Entry类型的parent，用来记录待存储元素的父结点。上面Entry的构造函数中，其中一个数据项就是父结点，

 

第15行-26行，当比较器不为空时，找出待存储元素的父结点，下面我们来逐行阅读

第18行，根据我们自定义的比较规则，待存储元素的键和当前结点比较大小，

第19行，如果小于当前结点的键，那么下一个比较的目标应在在左子树上，

第21行，如果大于当前结点的键，那么下一个比较的目标应该在右子树上，

第24行，如果待存储的键等于当前结点的键，那么将新值更新，同时返回旧值。

 

第27行-42行，当比较器为空时，采用默认的自然比较规则，找出父结点上述基本相似，

注意：第28行，对key是否为空校验，因此：TreeMap中的Key不允许为空。

 

第43行，赋值，

第44行-47行，确定新增元素时父结点的左结点还是右结点。

第49行-50行，size和modCount分别自增，

 

接下来我们着来分析第48行，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;
}

 夜，已深，至于fixAfterInsertion  how do it work ? 且听下回分解。