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的實現細節