集合-TreeMap源码解析

TreeMap是Map的一个实现类，底层数据结构是红黑树。它是有序的，默认顺序是自然排序，即ascii码升序，若在构造方法中传入比较器comparator，则按照比较器的规则进行排序。若对红黑树的数据结构不太熟悉，建议先移步https://blog.csdn.net/MAX_VALUE/article/details/104003631。本文主要介绍TreeMap的常用方法的工作流程，水平有限，若存在错误，请批评指正。

1.接口实现

NavigableMap<K,V>, Cloneable, Serializable

表示TreeMap是有序的，可拷贝的，可序列化的。

2.成员变量

//比较器
private final Comparator<? super K> comparator;

//红黑树根节点
private transient Entry<K,V> root;

//节点个数
private transient int size = 0;

//树结构被修改次数
private transient int modCount = 0;

//树节点
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;
}

排序比较器comparator(默认按key的ascii码升序)，Entry<K,V> root根节点，size键值对数量。

Entry就是一个红黑树节点，包含key,value,left,right,parent,color（left为左子树，ritht为右子树，parent为父节点，color为节点颜色）。

数据结构：红黑树

3.构造方法

    //无参构造    
    public TreeMap() {
        comparator = null;
    }

    //传入构造器
    public TreeMap(Comparator<? super K> comparator) {
        this.comparator = comparator;
    }

    //传入Map其他实现类，相当于做转换
    public TreeMap(Map<? extends K, ? extends V> m) {
        comparator = null;
        putAll(m);
    }

4.添加元素过程

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 {
                //从根节点开始，比较和插入节点的key值大小
                parent = t;
                cmp = k.compareTo(t.key);
                if (cmp < 0)
                    //将左子树设置为当前节点，继续查找
                    t = t.left;
                else if (cmp > 0)
                    //将右子树设置为当前节点，继续查找
                    t = t.right;
                else
                    //key值存在，更新value
                    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);
        //map中元素个数增加1，map结构被修改次数增加1
        size++;
        modCount++;
        return null;
    }

    /** From CLR */
    //对照红黑树的插入场景图会比较好理解，以下代码就是操作场景的翻译过程
    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
                    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;
    }

添加键值对，如果key值存在则替换旧的value值。

1.判断根节点是否存在，如果不存在将新增节点设置为根节点。结束

2.若根节点存在，判断是否传入比较器，如果有比较器根据比较器规则比较key值大小，如果未传入，则按照自然顺序进行排序（ascii码升序）。

3.找到插入节点的父节点。从根节点开始，比较插入节点的key值和当前节点的key值，如果大于将当前节点的右子树节点设置为当前节点，如果小于将当前节点的左子树节点设置为当前节点继续查找，如果等于则将该key值对应的value更新。循环判断直到当前节点为空，返回此时的父节点。

4.新增一个节点，该节点为红色。他的父节点是上述返回的父节点。

5.红黑树自平衡。场景具有对称性。3层判断（插入节点的父节点是否为红色、父节点是否是祖父节点的左子树、叔叔节点是否是红色、插入节点是父节点的左右子树）。

5.查找元素过程

根据key值查找指定的value

public V get(Object key) {
        //根据key查找entry树节点
        Entry<K,V> p = getEntry(key);
        return (p==null ? null : p.value);
    }

 

final Entry<K,V> getEntry(Object key) {
        // Offload comparator-based version for sake of performance
        if (comparator != null)
            //构造treeMap时，传入了比较器，则按比较器的规则排序进行查询
            //查找规则基于红黑树的查找算法
            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)
                //插入节点key小于当前节点key，向当前节点的左子树继续查找
                p = p.left;
            else if (cmp > 0)
                //插入节点key大于当前节点key，向当前节点的右子树继续查找
                p = p.right;
            else
                //插入节点key等于当前节点key，则p就是要查找的节点
                return p;
        }
        return null;
    }

1.传入比较器则按照比较器规则比较，否则按照自然顺序排序。

2.查找规则：将根节点设置为当前节点，若当前节点为空则直接返回null,若插入节点key值等于当前节点key值则返回当前节点，若小于则将当前节点的左子树设置为当前节点，若大于则将当前节点的右子树设置为当前节点。重复步骤2。

6.移除元素过程

public V remove(Object key) {
        //查找要删除的树节点p
        Entry<K,V> p = getEntry(key);
        if (p == null)
            return null;

        V oldValue = p.value;
        //删除树节点p
        deleteEntry(p);
        return oldValue;
    }

private void deleteEntry(Entry<K,V> p) {
        //map结构变化次数增加1，节点数减1
        modCount++;
        size--;

        // If strictly internal, copy successor's element to p and then make p
        // point to successor.
        //删除节点的两个子节点都存在
        if (p.left != null && p.right != null) {
            //查找下一个节点，在此处应该是删除节点p的右子树的最左子节点，即替换节点
            Entry<K,V> s = successor(p);
            //更新删除节点位置处的key,value为替换节点值
            p.key = s.key;
            p.value = s.value;
            //p指向替换节点
            p = s;
        } // p has 2 children

        
        //修正替换节点，删除节点只存在左子节点或只存在右子节点或不存在子节点
        //1.若只存在左节点，则替换节点就是左子节点
        //2.若只存在右节点，则替换节点就是右节点
        //3.不存在子节点则替换节点为空
        //4.删除节点的两个子节点都存在,若通过上述找到的替换节点存在右子节点，即替换节点
        //不是叶子节点，此时将替换节点更改为它的右子节点。
        //5.删除节点的两个子节点都存在,若通过上述找到的替换节点不存在右子节点,替代节点不变，
        //replacement 为null。
        
        // Start fixup at replacement node, if it exists.
        Entry<K,V> replacement = (p.left != null ? p.left : p.right);

        if (replacement != null) {
            // Link replacement to parent
            //场景1、2、4
            //解除删除节点引用方式1：使用替代节点替换删除节点的相关引用，删除节点的引用
            //置为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;

            // Null out links so they are OK to use by fixAfterDeletion.
            p.left = p.right = p.parent = null;

            // Fix replacement
            //替换节点是黑色节点，红黑树的自平衡
            if (p.color == BLACK)
                fixAfterDeletion(replacement);
        } else if (p.parent == null) { // return if we are the only node.
            //删除节点是根节点
            root = null;
        } else { //  No children. Use self as phantom replacement and unlink.
            //场景3、5
            if (p.color == BLACK)
                fixAfterDeletion(p);

            //解除删除节点引用方式2：删除节点的key,value值先替换成替换节点的key,value,
            //然后解除替代节点的引用关系
            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;
            }
        }
    }

    /** From CLR */
    //红色树删除操作场景算法
    private void fixAfterDeletion(Entry<K,V> x) {
        //x为替换节点，替换节点为黑色节点，做自平衡
        while (x != root && colorOf(x) == BLACK) {
            //替换节点是父节点的左子节点
            if (x == leftOf(parentOf(x))) {
                //替换接地的兄弟节点
                Entry<K,V> sib = rightOf(parentOf(x));

                //兄弟节点为红色
                if (colorOf(sib) == RED) {
                    //兄弟节点设置为黑色，父节点设置为红色，对父节点左旋
                    setColor(sib, BLACK);
                    setColor(parentOf(x), RED);
                    rotateLeft(parentOf(x));
                    //重新找到兄弟节点
                    sib = rightOf(parentOf(x));
                }
                
                //兄弟节点为黑色
                if (colorOf(leftOf(sib))  == BLACK &&
                    colorOf(rightOf(sib)) == BLACK) {
                    //兄弟节点的子节点都是黑色
                    //兄弟节点设置为红色，替换节点的父节点重新作为替换节点
                    setColor(sib, RED);
                    x = parentOf(x);
                } else {
                    if (colorOf(rightOf(sib)) == BLACK) {
                        //兄弟节点的右子节点是黑色，左子节点是红色
                        //设置兄弟节点为红色，兄弟节点的左子节点为黑色
                        //对兄弟节点进行右旋
                        setColor(leftOf(sib), BLACK);
                        setColor(sib, RED);
                        rotateRight(sib);
                        sib = rightOf(parentOf(x));
                    }
                    //兄弟节点的右子节点是红色，左子节点是任意色
                    //将兄弟节点颜色设置为父节点颜色，父节点设置为黑色，兄弟节点的
                    //右子节点设置为黑色，对父节点左旋
                    setColor(sib, colorOf(parentOf(x)));
                    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);
    }

根据key值移除键值对

1.根据key值查找到相应的节点，参考查找元素过程。

2.查找替代节点，若删除节点无子节点，则替代节点为本身；若删除节点有一个子节点，则替代节点为该子节点；若删除节点有两个子节点，则替代节点为右子树节点的最左节点。删除节点的方式有两种，方式1：使用替代节点替换删除节点的相关引用，删除节点的引用置为null。方式2：删除节点的key,value值先替换成替换节点的key,value,然后解除替代节点的引用关系。

3.红黑树完成自平衡。具有对称性。四层判断（替代节点是否为黑色，替代节点是否为其父节点的左子树，替代节点的兄弟节点是否为黑色，替代节点的兄弟节点的子树颜色）。确定好具体的场景进行变色和旋转，源码中的步骤同红黑树算法步骤。

7.TreeMap遍历

Map的遍历本质上还是基于Iterator的相关实现类。使用的是iterator()、next()方法。

for(Map.Entry<Integer,String> entry : map.entrySet()) {}

遍历步骤：

    public Set<Map.Entry<K,V>> entrySet() {
        //第一次调用entrySet()时，entrySet为null。
        EntrySet es = entrySet;
        //new EntrySet()会调用其iterator()方法
        return (es != null) ? es : (entrySet = new EntrySet());
    }


    class EntrySet extends AbstractSet<Map.Entry<K,V>> {
        public Iterator<Map.Entry<K,V>> iterator() {
            //查找红黑树的第一个节点，即最左子节点，遍历到的第一个元素就是该节点
            return new EntryIterator(getFirstEntry());
        }
    }

    public Map.Entry<K,V> next() {
        //之后调用iterator的next()方法
        return nextEntry();
    }

    final Entry<K,V> nextEntry() {
         Entry<K,V> e = next;
         if (e == null)
            throw new NoSuchElementException();
         if (modCount != expectedModCount)
            throw new ConcurrentModificationException();
         //按照红黑树的排序顺序找到下一个节点
         next = successor(e);
         lastReturned = e;
         return e;
    }

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 {
            //1.若当前节点的是其父节点的右子节点，继续查找其父节点，        
            //下一个节点是其祖父节点或者null
            //2.若当前节点的是其父节点的左子节点，则下一个节点是其父节点
            Entry<K,V> p = t.parent;
            Entry<K,V> ch = t;
            while (p != null && ch == p.right) {
                ch = p;
                p = p.parent;
            }
            return p;
        }
    }

8.TreeSet源码分析

TreeSet<E>的底层是NavigableMap<E,Object>，本质是TreeMap<E,Object>。TreeSet的增删改查是基于TreeMap的增删改查来实现的。

//m是TreeSet的底层数据结构
private transient NavigableMap<E,Object> m;    

//通常使用无参构造
public TreeSet() {
   this(new TreeMap<E,Object>());
}

//默认将TreeMap作为TreeSet底层的实现
TreeSet(NavigableMap<E,Object> m) {
   this.m = m;
}