TreeMap源码详解(jdk1.8)

[diaosichengxuyuan]

1.准备工作
TreeMap的实现是红黑树，所以请先熟悉本人原创的红黑树原理一文：diaosichengxuyuan/DataStructureAlgorithm#4
本文中源码的解释就是基于此文中的图示，方便大家理解。另外，TreeSet可以实现一致性哈希算法

2.基础数据模型

public class TreeMap<K,V> extends AbstractMap<K,V> implements NavigableMap<K,V>, Cloneable, java.io.Serializable
{
   //可以构造Tree的时候自定义构造器Comparator，也可以让key实现Comparable接口，也就是说key一定
   //是可排序的
   private final Comparator<? super K> comparator;
   
   //树节点
   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;
        //黑色为true，红色为false
        boolean color = BLACK;
   }
}

3.查找

    public V get(Object key) {
        Entry<K,V> p = getEntry(key);
        return (p==null ? null : p.value);
    }
    
    final Entry<K,V> getEntry(Object key) {
        //使用Comparator比较器
        if (comparator != null)
            return getEntryUsingComparator(key);
        //不支持key为null的场景
        if (key == null)
            throw new NullPointerException();
       
        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;
    }
    
    //使用Comparator比较器，逻辑跟上面是一样的
    final Entry<K,V> getEntryUsingComparator(Object key) {
        @SuppressWarnings("unchecked")
            K k = (K) key;
        Comparator<? super K> cpr = comparator;
        if (cpr != null) {
            Entry<K,V> p = root;
            while (p != null) {
                int cmp = cpr.compare(k, p.key);
                if (cmp < 0)
                    p = p.left;
                else if (cmp > 0)
                    p = p.right;
                else
                    return p;
            }
        }
        return null;
    }

4.插入
首先我们看下左旋转和右旋转的代码

    //可以对照原理一文中左旋转的图进行理解，此处是对节点p左旋转，我们的图中是对节点A左旋转，所以
    //我说的A就是现在的节点p
    private void rotateLeft(Entry<K,V> p) {
        if (p != null) {
            //记住B
            Entry<K,V> r = p.right;
            //B的左孩子(黄色节点)变成A的右孩子
            p.right = r.left;
            //更新黄色节点的父亲是A
            if (r.left != null)
                r.left.parent = p;
            //更新B的父亲
            r.parent = p.parent;
            //B作为A父亲的孩子
            if (p.parent == null)
                root = r;
            else if (p.parent.left == p)
                p.parent.left = r;
            else
                p.parent.right = r;
            //B的左孩子是A
            r.left = p;
            //A的父亲是B
            p.parent = r;
        }
    }

    //跟左旋转相反，不多解释
    private void rotateRight(Entry<K,V> p) {
        if (p != null) {
            Entry<K,V> l = p.left;
            p.left = l.right;
            if (l.right != null) l.right.parent = p;
            l.parent = p.parent;
            if (p.parent == null)
                root = l;
            else if (p.parent.right == p)
                p.parent.right = l;
            else p.parent.left = l;
            l.right = p;
            p.parent = l;
        }
    }

一些其他的方法，大体看一下

private static <K,V> boolean colorOf(Entry<K,V> p) {
        return (p == null ? BLACK : p.color);
    }

    private static <K,V> Entry<K,V> parentOf(Entry<K,V> p) {
        return (p == null ? null: p.parent);
    }

    private static <K,V> void setColor(Entry<K,V> p, boolean c) {
        if (p != null)
            p.color = c;
    }

    private static <K,V> Entry<K,V> leftOf(Entry<K,V> p) {
        return (p == null) ? null: p.left;
    }

    private static <K,V> Entry<K,V> rightOf(Entry<K,V> p) {
        return (p == null) ? null: p.right;
    }

插入的逻辑

public V put(K key, V value) {
        Entry<K,V> t = root;
        //如果是一颗空树
        if (t == null) {
            //提前把key=null或者key不能排序的异常抛出来
            compare(key, key); 
            //组成新的根节点
            root = new Entry<>(key, value, null);
            size = 1;
            modCount++;
            return null;
        }
        
        int cmp;
        Entry<K,V> parent;
        Comparator<? super K> cpr = comparator;
        //使用传入的Comparator比较器
        if (cpr != null) {
            do {
                //使用while循环而不使用递归，每次要把父亲节点提前记住，如果不记住最后t=null退出，不知道t
                //的父节点是谁
                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);
        }
        //使用Comparable比较器
        else {
            if (key == null)
                throw new NullPointerException();            
                Comparable<? super K> k = (Comparable<? super  K>) key;
            do {
                //使用while循环而不使用递归，每次要把父亲节点提前记住，如果不记住最后t=null退出，不知道t
                //的父节点是谁
                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);
        }
        //添加一个新的节点e，其父亲是parent
        Entry<K,V> e = new Entry<>(key, value, parent);
        //e是parent左孩子
        if (cmp < 0)
            parent.left = e;
        //e是parent右孩子
        else
            parent.right = e;
        //调整红黑树结构
        fixAfterInsertion(e);
        //元素个数加1
        size++;
        //变更次数加1
        modCount++;
        return null;
    }

上面代码中我觉得有两点体现了工匠精神：1.循环使用的是do while而不是while，因为已经判断了t!=null，所以do while的第一次肯定是成立的，比while循环减少一次判断。2.使用cmp在迭代过程中直接确定了e是parent的左孩子还是右孩子，不用在添加的时候多一次计算。

其实红黑树插入过程与二叉排序树的插入过程一样，不同点在于插入完成后需要执行调整方法fixAfterInsertion，下面看下调整过程：

//我们对照原理一文中插入的6种情况以及配图讲解
private void fixAfterInsertion(Entry<K,V> x) {
        //插入的节点肯定是红色的
        x.color = RED;

        //对于情况2，x.parent.color == BLACK，不会进入下面的while
        //循环，因为不需要调整

        //N的父亲P必须是红色的
        while (x != null && x != root && x.parent.color == RED) {
            //P是G的左孩子
            if (parentOf(x) == leftOf(parentOf(parentOf(x))))  {               
                Entry<K,V> y =rightOf(parentOf(parentOf(x)));
                //情况3，叔叔A为红色
                if (colorOf(y) == RED) {
                    setColor(parentOf(x), BLACK);
                    setColor(y, BLACK);
                    setColor(parentOf(parentOf(x)), RED);
                    x = parentOf(parentOf(x));
                } else {
                    //情况5，N是P的右孩子
                    if (x == rightOf(parentOf(x))) {
                        x = parentOf(x);
                        rotateLeft(x);
                    }
                    //情况4，N是P的左孩子
                    setColor(parentOf(x), BLACK);
                    setColor(parentOf(parentOf(x)), RED);
                    rotateRight(parentOf(parentOf(x)));
                }
            }
            //情况6，P是G的右孩子，只需要把上面的旋转方向变一下就行了
            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)));
                }
            }
        }
        //情况1，根节点涂黑即可
        root.color = BLACK;
    }

5.删除

    public V remove(Object key) {
        //先查找需要删除的元素
        Entry<K,V> p = getEntry(key);
        //找不到直接返回
        if (p == null)
            return null;

        V oldValue = p.value;
        //删除
        deleteEntry(p);
        return oldValue;
    }
    
    private void deleteEntry(Entry<K,V> p) {
        //变更次数加1
        modCount++;
        //元素个数减1
        size--;

        //情况3，可以转为情况2处理
        if (p.left != null && p.right != null) {
            //查找用来替换p的节点，这里使用的是p右子树最左的节点，我们在原理一文中使用的是左子树最右
            //的节点，两种方式都可以。其实s肯定没有左孩子，它最多有右孩子或者没有孩子。
            Entry<K,V> s = successor(p);
            //直接用s替换掉p，但是p的引用关系以及颜色不改变
            p.key = s.key;
            p.value = s.value;
            //将s作为新的要删除的节点，此时已经转成情况2了
            p = s;
        }

        //按照上面说的，p.left肯定是null，p.right可能为null也可能
        //不为null
        Entry<K,V> replacement = (p.left != null ? p.left : p.right);
        
        //p.right!=null时
        if (replacement != null) {
            //删除p，使用replacement替代
            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.color==RED，属于情况4，不需要调整
            //若p.color== BLACK，需要调整
            if (p.color == BLACK)
                fixAfterDeletion(replacement);
        }
        //p.right==null时，说明p是一个叶子节点，属于情况1
        //如果该叶子节点是根
        else if (p.parent == null) {
            root = null;
        }
        //如果该叶子节点不是根
        else {
            //情况1中，N是黑色时，需要调整
            if (p.color == BLACK)
                fixAfterDeletion(p);
            //情况1中，N是红色时，直接删除
            if (p.parent != null) {
                if (p == p.parent.left)
                    p.parent.left = null;
                else if (p == p.parent.right)
                    p.parent.right = null;
                p.parent = null;
            }
        }
    }
    
    //找替换t的节点，使用的是t右子树上最左的节点
    static <K,V> TreeMap.Entry<K,V> successor(Entry<K,V> t) {
        if (t == null)
            return null;
        else if (t.right != null) {
            Entry<K,V> p = t.right;
            while (p.left != null)
                p = p.left;
            return p;
        } else {
            Entry<K,V> p = t.parent;
            Entry<K,V> ch = t;
            while (p != null && ch == p.right) {
                ch = p;
                p = p.parent;
            }
            return p;
        }
    }

删除后的调整过程

    //看这个代码之前，要先想明白，要把x当成replacement，也就是x现在是已经将该删除的节点删除后的
    //替换节点，比如在情况6的图中现在x不是N，而是已经将N替换掉的NL，NL颜色为红，只是图不大匹
    //配，所以对于情况6，while循环直接不成立就跳到最后一行将NL染黑退出，这也正是我们情况6的处理 
    //方式
    private void fixAfterDeletion(Entry<K,V> x) {
        //x为黑时，对应的是7 8 9 10 11的场景，配图中的解释可能与代码有些出入，我们可以把图中的N当做 
        //是已经将节点删除后重新替换的节点，因为此时NL也是黑色的，把N当成NL不存在问题。图中的N就 
        //是现在的入参x，此时P的左路径黑色节点数已经比右路径黑色节点数少1了
        while (x != root && colorOf(x) == BLACK) {
            //N是P的左孩子
            if (x == leftOf(parentOf(x))) {
                //叔叔节点B
                Entry<K,V> sib = rightOf(parentOf(x));
                //情况7，叔叔为红色，处理方法就是将其转化成叔叔为黑色的情况，所以经过这段代码处理叔叔 
                //都是黑色的
                if (colorOf(sib) == RED) {
                    setColor(sib, BLACK);
                    setColor(parentOf(x), RED);
                    rotateLeft(parentOf(x));
                    sib = rightOf(parentOf(x));
                }
                
                //情况8和9，这块代码对于情况8比较好理解，将B染红，然后将P作为需要处理的节点继续向上处 
                //理，但是情况9中将B染红后，为什么没有将P染黑呢？其实虽然此时P是黑色的，下次循环while 
                //不成立，跳转到代码最后就会将P染黑了
                if (colorOf(leftOf(sib))  == BLACK &&
                    colorOf(rightOf(sib)) == BLACK) {
                    setColor(sib, RED);
                    x = parentOf(x);
                }
                //情况10和11
                else {
                    //情况11，右孩子为黑，其实左孩子肯定为红，为什么？因为如果左孩子为黑，就是情况8和9， 
                    //我们上面已经处理过了，所以情况11转成了情况10处理
                    if (colorOf(rightOf(sib)) == BLACK) {
                        setColor(leftOf(sib), BLACK);
                        setColor(sib, RED);
                        rotateRight(sib);
                        sib = rightOf(parentOf(x));
                    }
                    //情况10
                    setColor(sib, colorOf(parentOf(x)));
                    setColor(parentOf(x), BLACK);
                    setColor(rightOf(sib), BLACK);
                    rotateLeft(parentOf(x));
                    //为什么将root赋值给x而不是直接break跳出呢？看下情况10的图，此时的顶点变成了B，而B是 
                    //不确定颜色的，如果B是根，我们要保证它是黑色的
                    x = root;
                }
            }
            //N是P的右孩子，左右旋转反一下就可以了，不再详细解释
            else {
                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;
                }
            }
        }
        
        //情况5和情况6，将NL染黑。不过仔细想一下，情况5应该不可能，当前节点是替换节点，有哪个节点 
        //的replacement会是根呢？
        setColor(x, BLACK);
    } 

其实还有一点非常有意思，deleteEntry方法中当p是叶子节点时，是先调用fixAfterDeletion调整，然后再将p删除，而叶子节点删除时，是先将p删除，用replacement替换p，然后再调整。想一想为什么呢？其实非叶子节点replacement替换p之后，replacement的上下引用关系跟p原来是一模一样的，fixAfterDeletion方法逻辑正确，而叶子节点如果先删除，那它的replacement可就是null了，fixAfterDeletion就没法用了。至于到底应该先删再调整还是先调整再删，其实都可以，只是代码就不能这么写了。

作者原创，转载请注明出处，违法必究！