java 集合框架-TreeMap

一、背景

1、SortMap接口

• 扩展Map接口，定义按照key有序的映射集合
• 以key的自然顺序（实现Comparable的对象）排序或者指定的Comparator排序
• 有序指的是迭代的有序，如entrySet、keySet、values 等方法返回的元素集合有序

SortMap新增定义了一些基于有序的方法：

//返回一个左闭右开区间的子视图，修改子视图等同修改该map
SortedMap<K,V> subMap(K fromKey, K toKey);

//返回小于指定key（toKey）的子视图，修改子视图等同修改该map
SortedMap<K,V> headMap(K toKey);

//返回大于或等于指定key（fromKey）的子视图，修改子视图等同修改该map
SortedMap<K,V> tailMap(K fromKey);

//返回该map第一个key（最小的）
K firstKey();

//返回该map最后一个key（最大的）
K lastKey();

2、NavigableMap 接口

• 扩展SortMap，增加定义一组可以搜索指定key的最接近键值（排序）
• 定义可以正向或者反向顺序访问的映射集合

NavigableMap新增定义了一些方法：

//返回key小于指定key的最大的的键值对，如没有这样的元素返回null
Map.Entry<K,V> lowerEntry(K key);

//返回key小于指定key的最大key，如没有这个样key返回null
K lowerKey(K key);

//返回key小于或等于指定key的最大的的键值对，如没有这样的元素返回null
Map.Entry<K,V> floorEntry(K key);

//返回key小于或等于指定key的最大key，如没有这个样key返回null
K floorKey(K key);

//返回大于或等于指定key的最小键值对，如没有这样的元素返回null
Map.Entry<K,V> ceilingEntry(K key);

//返回大于或等于指定key的最小key，如没有这样的元素返回null
K ceilingKey(K key);

//返回大于指定key的最小键值对，如没有这样的元素返回null
Map.Entry<K,V> higherEntry(K key);

//返回大于指定key的最小key，如没有这样的元素返回null
K higherKey(K key);

//返回key最小的键值对，如没有这样的元素返回null
Map.Entry<K,V> firstEntry();

//返回key最大的键值对，如没有这样的元素返回null
Map.Entry<K,V> lastEntry();

//返回并删除最小的键值对，如没有这样的元素返回null
Map.Entry<K,V> pollFirstEntry();

//返回并删除key最大的键值对，如没有这样的元素返回null
Map.Entry<K,V> pollLastEntry();

//返回一个反序的NavigableMap视图
NavigableMap<K,V> descendingMap();

//获取key的 NavigableSet 视图
NavigableSet<K> navigableKeySet();

//获取key反序的 NavigableSet 视图
NavigableSet<K> descendingKeySet();

//扩展SortMap中的subMap，支持指定区间的开或闭
NavigableMap<K,V> subMap(K fromKey, boolean fromInclusive,
                             K toKey,   boolean toInclusive);

//返回小于或等于（inclusive指定）指定key（toKey）的子视图                      
NavigableMap<K,V> headMap(K toKey, boolean inclusive);

//返回大于或等于（inclusive指定）指定key（fromKey）的子视图
NavigableMap<K,V> tailMap(K fromKey, boolean inclusive);

二、概述

• 基于红黑树实现NavigableMap接口，实现以key自然有序或者指定Comparator有序
• 对于containsKey、get、put、remove方法时间复杂度为log(n)
• 非线程安全，并发需要自实现同步或者通过Collections.synchronizedSortedMap封装
• 与其他非线程安全集合实现类一样，支持迭代快速失败机制

要理解TreeMap的原来，必须要先清楚红黑树这种数据结构

1、红黑树

平衡二叉查找树：它是一棵空树或者它的左右两个子树的高度差绝对值不超过1，并且左右两个子树也是一棵平衡二叉查找树

红黑树是一种自平衡二叉查找树，具有以下特性：

• 每个节点是黑色或者红色
• 根节点是黑色
• 每个叶子节点（为空NIL或NULL）是黑色
• 如果一个节点是红色，则它的子节点必须是黑色
• 从一个节点到该节点的子孙节点的所有路径上包含相同数目的黑节点

二、成员变量

/**
* 用于key排序的比较器，构造方法传入，如果使用自然排序，则为空；<br>
* 后续的实现中使用那种排序就是通过判断这个变量是否为空
*/
private final Comparator<? super K> comparator;

/**
* 红黑树的根节点，初始为空树
*/
private transient Entry<K,V> root = null;

/**
* 记录元素（键值对）数量
*/
private transient int size = 0;

/**
* 修改次数，用于迭代快速失败机制
*/
private transient int modCount = 0;

三、构造方法

    public TreeMap() {
        comparator = null;
    }

    public TreeMap(Comparator<? super K> comparator) {
        this.comparator = comparator;
    }

    public TreeMap(Map<? extends K, ? extends V> m) {
        comparator = null;
        putAll(m);
    }
    
    public TreeMap(SortedMap<K, ? extends V> m) {
        comparator = m.comparator();
        try {
            buildFromSorted(m.size(), m.entrySet().iterator(), null, null);
        } catch (java.io.IOException cannotHappen) {
        } catch (ClassNotFoundException cannotHappen) {
        }
    }

提供4个构造方法，这4个构造方法也是SortMap接口中建议的4个

四、核心代码

先来看下键值对Entry定义：

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

        /**
         * Make a new cell with given key, value, and parent, and with
         * {@code null} child links, and BLACK color.
         */
        Entry(K key, V value, Entry<K,V> parent) {
            this.key = key;
            this.value = value;
            this.parent = parent;
        }

        /**
         * Returns the key.
         *
         * @return the key
         */
        public K getKey() {
            return key;
        }

        /**
         * Returns the value associated with the key.
         *
         * @return the value associated with the key
         */
        public V getValue() {
            return value;
        }

        /**
         * Replaces the value currently associated with the key with the given
         * value.
         *
         * @return the value associated with the key before this method was
         *         called
         */
        public V setValue(V value) {
            V oldValue = this.value;
            this.value = value;
            return oldValue;
        }

        public boolean equals(Object o) {
            if (!(o instanceof Map.Entry))
                return false;
            Map.Entry<?,?> e = (Map.Entry<?,?>)o;

            return valEquals(key,e.getKey()) && valEquals(value,e.getValue());
        }

        public int hashCode() {
            int keyHash = (key==null ? 0 : key.hashCode());
            int valueHash = (value==null ? 0 : value.hashCode());
            return keyHash ^ valueHash;
        }

        public String toString() {
            return key + "=" + value;
        }
    }

可以看出，Entry就是红黑树的节点，除了成员变量其他方法基本和其他Map的Entry一样：多了左右子节点的指针，指向父节点的指针以及表示颜色的boolean变量color；

下面主要看下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();
            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;
    }

首先，对空树（初始）有单独判断的处理，只需要将新增的元素作为根节点就完成了，这里有需要注意的是做了一次该新增key的compare操作，目的是为了做类型校验，保证key是实现了Comparable接口或者能使用Comparator比较；

    final int compare(Object k1, Object k2) {
        return comparator==null ? ((Comparable<? super K>)k1).compareTo((K)k2)
            : comparator.compare((K)k1, (K)k2);
    }

如果该map使用自然排序（构造时comparator为空），则强制转换key为Comparable进行比较如果类型不对会抛出ClassCastException，然put操作失败；如果使用指定Comparator也会强制转换为这个比较器的泛型，同样，如果类型不对，也会抛出ClassCastException

然后，如果不是空树，其实就是进行一般二叉查找树的插入操作，循环判断应该要插入的位置（红黑树的插入过程没有特殊）；在完成插入后，需要进行树的调整（旋转、着色）：

    /** 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 = 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;
    }

这个方法就是红黑树插入新节点后的调整算法，具体算法这里不做描述，可以参考具体的算法文章或书籍

接下来看下get方法的实现：

    public V get(Object key) {
        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)
            return getEntryUsingComparator(key);
        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;
    }

TreeMap中，包括get、contains、containsKey、remove等方法都是通过getEntry来获取键值对；getEntry其实就是二叉查找树的查找算法，通过大小判断循环在左或右字数查找，这里还特殊处理comparator不为空的情况下使用getEntryUsingComparator方法查找，其实实现算法都是一样的，只是两个key比较的方式不一样，我个人觉得完全是不需要单独一个方法的，只需要将两个key的比较使用统一的compare方法就可以

接下来看下删除remove方法：

    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) {
        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) {
            Entry<K,V> s = successor(p);
            p.key = s.key;
            p.value = s.value;
            p = s;
        } // p has 2 children

        // 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
            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.
            if (p.color == BLACK)
                fixAfterDeletion(p);

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

首先通过getEntry查找到需要删除的节点，然后使用deleteEntry删除指定节点，deleteEntry的实现就是通过断接节点的指针来删除指定节点，需要判断的情况多；同样删除后调整红黑树，具体的fixAfterDeletion方法就不贴了

基于二叉查找树，SortMap接口定义的方法就非常容易实现了，比如

    /**
     * @throws NoSuchElementException {@inheritDoc}
     */
    public K firstKey() {
        return key(getFirstEntry());
    }

    final Entry<K,V> getFirstEntry() {
        Entry<K,V> p = root;
        if (p != null)
            while (p.left != null)
                p = p.left;
        return p;
    }

获取排序第一个的元素，获取二叉查找树最左的节点即可，同理lastKey获取最右的节点，其他可以基于比较排序的方法也是一样，如lowerEntry、ceilingEntry等，这里不一一贴出

对于subMap这类获取子视图的方法（headMap、tailMap等），TreeMap实现了一个通用的抽象内部类实现：

    abstract static class NavigableSubMap<K,V> extends AbstractMap<K,V>
        implements NavigableMap<K,V>, java.io.Serializable {
        /**
         * The backing map.
         */
        final TreeMap<K,V> m;

        /**
         * Endpoints are represented as triples (fromStart, lo,
         * loInclusive) and (toEnd, hi, hiInclusive). If fromStart is
         * true, then the low (absolute) bound is the start of the
         * backing map, and the other values are ignored. Otherwise,
         * if loInclusive is true, lo is the inclusive bound, else lo
         * is the exclusive bound. Similarly for the upper bound.
         */
        final K lo, hi;
        final boolean fromStart, toEnd;
        final boolean loInclusive, hiInclusive;

        NavigableSubMap(TreeMap<K,V> m,
                        boolean fromStart, K lo, boolean loInclusive,
                        boolean toEnd,     K hi, boolean hiInclusive) {
            if (!fromStart && !toEnd) {
                if (m.compare(lo, hi) > 0)
                    throw new IllegalArgumentException("fromKey > toKey");
            } else {
                if (!fromStart) // type check
                    m.compare(lo, lo);
                if (!toEnd)
                    m.compare(hi, hi);
            }

            this.m = m;
            this.fromStart = fromStart;
            this.lo = lo;
            this.loInclusive = loInclusive;
            this.toEnd = toEnd;
            this.hi = hi;
            this.hiInclusive = hiInclusive;
        }
    }

这个类的代码有点长，这里只贴了成员变量和构造方法，只说下实现的原理，成员变量 m 保持了原本TreeMap的引用，所以子map的修改其实是直接修改本身map的；其他几个成员变量分别表示：lo、hi：表示区间的两端key值；fromStart、toEnd：是否有头尾区间，false表示没有固定开始/结束key，开始/结束就是所有元素的头/尾；loInclusive、 hiInclusive：表示区间的开或闭；
具体的成员方法实现都是大同小异，基本每个操作都会先判断入参的key是否在这个子Map的范围内，比如：

        public final boolean containsKey(Object key) {
            return inRange(key) && m.containsKey(key);
        }

        public final V put(K key, V value) {
            if (!inRange(key))
                throw new IllegalArgumentException("key out of range");
            return m.put(key, value);
        }

最后来看迭代器，和其他Map实现一样，有EntrySet、keySet、values 三个迭代器，也是同样实现了一个抽象迭代器器，具体迭代器实现个抽象类

    abstract class PrivateEntryIterator<T> implements Iterator<T> {
        Entry<K,V> next;
        Entry<K,V> lastReturned;
        int expectedModCount;

        PrivateEntryIterator(Entry<K,V> first) {
            expectedModCount = modCount;
            lastReturned = null;
            next = first;
        }

        public final boolean hasNext() {
            return next != null;
        }

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

        final Entry<K,V> prevEntry() {
            Entry<K,V> e = next;
            if (e == null)
                throw new NoSuchElementException();
            if (modCount != expectedModCount)
                throw new ConcurrentModificationException();
            next = predecessor(e);
            lastReturned = e;
            return e;
        }

        public void remove() {
            if (lastReturned == null)
                throw new IllegalStateException();
            if (modCount != expectedModCount)
                throw new ConcurrentModificationException();
            // deleted entries are replaced by their successors
            if (lastReturned.left != null && lastReturned.right != null)
                next = lastReturned;
            deleteEntry(lastReturned);
            expectedModCount = modCount;
            lastReturned = null;
        }
    }

可以看到，这完全就是对一棵二叉查询树的遍历算法，successor 方法是找节点的顺序下一个节点

五、总结

总体来说，TreeMap的原理完全就是红黑树，要理解红黑树的算法才能清楚TreeMap的实现细节