深入源碼分析TreeSet和TreeMap

類似於前面介紹的HashMap和HashSet之間的關系，HashSet底層依賴於HashMap實現，而TreeSet底層則采用一個NavigableMap來保存TreeSet集合的元素。但實際上，由於NavigableMap只是一個接口，因底層依然是使用TreeMap來包含Set集合中的所有元素。

我們觀察TreeSet的部分源碼:

package java.util;

public class TreeSet<E> extends AbstractSet<E>
    implements NavigableSet<E>, Cloneable, java.io.Serializable
{
    // 使用NavigableMap對象的key來保存Set集合的元素
    private transient NavigableMap<E,Object> m;

    //使用PRESENT作為Map集合中的value
    private static final Object PRESENT = new Object();

    // 不帶參數的構造函數。創建一個空的TreeMap
    //以自然排序方法創建一個新的TreeMap，再根據該TreeMap創建一個TreeSet
    //使用該TreeMap的key來保存Set集合的元素
    public TreeSet() {
        this(new TreeMap<E,Object>());
    }

    // 將TreeMap賦值給 "NavigableMap對象m"
    TreeSet(NavigableMap<E,Object> m) {
        this.m = m;
    }

    //以定制排序的方式創建一個新的TreeMap。根據該TreeMap創建一個TreeSet
    //使用該TreeMap的key來保存set集合的元素
    public TreeSet(Comparator<? super E> comparator) {
        this(new TreeMap<E,Object>(comparator));
    }

    // 創建TreeSet，並將集合c中的全部元素都添加到TreeSet中
    public TreeSet(Collection<? extends E> c) {
        this();
        // 將集合c中的元素全部添加到TreeSet中
        addAll(c);
    }

    // 創建TreeSet，並將s中的全部元素都添加到TreeSet中
    public TreeSet(SortedSet<E> s) {
        this(s.comparator());
        addAll(s);
    }

    // 返回TreeSet的順序排列的迭代器。
    // 因為TreeSet時TreeMap實現的，所以這裡實際上時返回TreeMap的“鍵集”對應的迭代器
    public Iterator<E> iterator() {
        return m.navigableKeySet().iterator();
    }

    // 返回TreeSet的逆序排列的迭代器。
    // 因為TreeSet時TreeMap實現的，所以這裡實際上時返回TreeMap的“鍵集”對應的迭代器
    public Iterator<E> descendingIterator() {
        return m.descendingKeySet().iterator();
    }

    // 返回TreeSet的大小
    public int size() {
        return m.size();
    }

    // 返回TreeSet是否為空
    public boolean isEmpty() {
        return m.isEmpty();
    }

    // 返回TreeSet是否包含對象(o)
    public boolean contains(Object o) {
        return m.containsKey(o);
    }

    // 添加e到TreeSet中
    public boolean add(E e) {
        return m.put(e, PRESENT)==null;
    }

    // 刪除TreeSet中的對象o
    public boolean remove(Object o) {
        return m.remove(o)==PRESENT;
    }

    // 清空TreeSet
    public void clear() {
        m.clear();
    }

    // 將集合c中的全部元素添加到TreeSet中
    public  boolean addAll(Collection<? extends E> c) {
        // Use linear-time version if applicable
        if (m.size()==0 && c.size() > 0 &&
            c instanceof SortedSet &&
            m instanceof TreeMap) {
            //把C集合強制轉換為SortedSet集合
            SortedSet<? extends E> set = (SortedSet<? extends E>) c; 
             //把m集合強制轉換為TreeMap集合
            TreeMap<E,Object> map = (TreeMap<E, Object>) m;
            Comparator<? super E> cc = (Comparator<? super E>) set.comparator();
            Comparator<? super E> mc = map.comparator();
            //如果cc和mc兩個Comparator相等
            if (cc==mc || (cc != null && cc.equals(mc))) {
            //把Collection中所有元素添加成TreeMap集合的key
                map.addAllForTreeSet(set, PRESENT);
                return true;
            }
        }
        return super.addAll(c);
    }

    // 返回子Set，實際上是通過TreeMap的subMap()實現的。
    public NavigableSet<E> subSet(E fromElement, boolean fromInclusive,
                                  E toElement,   boolean toInclusive) {
        return new TreeSet<E>(m.subMap(fromElement, fromInclusive,
                                       toElement,   toInclusive));
    }

    // 返回Set的頭部，范圍是：從頭部到toElement。
    // inclusive是是否包含toElement的標志
    public NavigableSet<E> headSet(E toElement, boolean inclusive) {
        return new TreeSet<E>(m.headMap(toElement, inclusive));
    }

    // 返回Set的尾部，范圍是：從fromElement到結尾。
    // inclusive是是否包含fromElement的標志
    public NavigableSet<E> tailSet(E fromElement, boolean inclusive) {
        return new TreeSet<E>(m.tailMap(fromElement, inclusive));
    }

    // 返回子Set。范圍是：從fromElement(包括)到toElement(不包括)。
    public SortedSet<E> subSet(E fromElement, E toElement) {
        return subSet(fromElement, true, toElement, false);
    }

    // 返回Set的頭部，范圍是：從頭部到toElement(不包括)。
    public SortedSet<E> headSet(E toElement) {
        return headSet(toElement, false);
    }

    // 返回Set的尾部，范圍是：從fromElement到結尾(不包括)。
    public SortedSet<E> tailSet(E fromElement) {
        return tailSet(fromElement, true);
    }

    // 返回Set的比較器
    public Comparator<? super E> comparator() {
        return m.comparator();
    }

    // 返回Set的第一個元素
    public E first() {
        return m.firstKey();
    }

    // 返回Set的最後一個元素
    public E first() {
    public E last() {
        return m.lastKey();
    }

    // 返回Set中小於e的最大元素
    public E lower(E e) {
        return m.lowerKey(e);
    }

    // 返回Set中小於/等於e的最大元素
    public E floor(E e) {
        return m.floorKey(e);
    }

    // 返回Set中大於/等於e的最小元素
    public E ceiling(E e) {
        return m.ceilingKey(e);
    }

    // 返回Set中大於e的最小元素
    public E higher(E e) {
        return m.higherKey(e);
    }

    // 獲取第一個元素，並將該元素從TreeMap中刪除。
    public E pollFirst() {
        Map.Entry<E,?> e = m.pollFirstEntry();
        return (e == null)? null : e.getKey();
    }

    // 獲取最後一個元素，並將該元素從TreeMap中刪除。
    public E pollLast() {
        Map.Entry<E,?> e = m.pollLastEntry();
        return (e == null)? null : e.getKey();
    }

    // 克隆一個TreeSet，並返回Object對象
    public Object clone() {
        TreeSet<E> clone = null;
        try {
            clone = (TreeSet<E>) super.clone();
        } catch (CloneNotSupportedException e) {
            throw new InternalError();
        }

        clone.m = new TreeMap<E,Object>(m);
        return clone;
    }

    // java.io.Serializable的寫入函數
    // 將TreeSet的“比較器、容量，所有的元素值”都寫入到輸出流中
    private void writeObject(java.io.ObjectOutputStream s)
        throws java.io.IOException {
        s.defaultWriteObject();

        // 寫入比較器
        s.writeObject(m.comparator());

        // 寫入容量
        s.writeInt(m.size());

        // 寫入“TreeSet中的每一個元素”
        for (Iterator i=m.keySet().iterator(); i.hasNext(); )
            s.writeObject(i.next());
    }

    // java.io.Serializable的讀取函數：根據寫入方式讀出
    // 先將TreeSet的“比較器、容量、所有的元素值”依次讀出
    private void readObject(java.io.ObjectInputStream s)
        throws java.io.IOException, ClassNotFoundException {
        // Read in any hidden stuff
        s.defaultReadObject();

        // 從輸入流中讀取TreeSet的“比較器”
        Comparator<? super E> c = (Comparator<? super E>) s.readObject();

        TreeMap<E,Object> tm;
        if (c==null)
            tm = new TreeMap<E,Object>();
        else
            tm = new TreeMap<E,Object>(c);
        m = tm;

        // 從輸入流中讀取TreeSet的“容量”
        int size = s.readInt();

        // 從輸入流中讀取TreeSet的“全部元素”
        tm.readTreeSet(size, s, PRESENT);
    }

    // TreeSet的序列版本號
    private static final long serialVersionUID = -2479143000061671589L;
}

從上述可以看出，TreeSet的構造函數都是通過新建一個TreeMap作為實際存儲Set元素的容器。因此得出結論: TreeSet的底層實際使用的存儲容器就是TreeMap。

對於TreeMap而言，它采用一種被稱為”紅黑樹”的排序二叉樹來保存Map中每個Entry。每個Entry被當成”紅黑數”的一個節點來對待。

對於如下代碼:

TreeMap<String,Double> map = new TreeMap<String,Double>();
map.put("ccc",89.0);
map.put("aaa",80.0);
map.put("bbb",89.0);
map.put("bbb",89.0);

如上代碼所示，程序向TreeMap放入四個值。根據”紅黑數”的定義，程序會把”ccc,89.0”這個Entry做為該”紅黑數”的根節點，然後執行put(“aaa”,”80.0”)時會將其作為新的節點添加到已存在的紅黑樹中。也就是說，以後每次向TreeMap中放入一個key-value對，系統都需要將Entry當成一個新的節點，添加到存在的”紅黑數”中，通過這種方式就可以保證TreeMap中所有的key都是按一定順序地排列的。

我們簡單總結一下HashMap，HashSet與TreeMap，TreeSet兩類集合。

• HashMap，HashSet的存儲集合元素的方式為不同的東西放在不同的位置，需要時就可以快速的找到它們。
• TreeMap,TreeSet的存儲方式為根據一個排序規則來為所有的元素排序，然後該集合存儲的就是排序好的元素。

由於TreeMap底層采用一顆”紅黑數”來保存集合中的Entry。所以TreeMap添加元素，取出元素的性能都比HashMap低。當TreeMap添加元素時，需要通過循環找到新增的Entry的插入位置，因為比較耗性能。當取出元素時，也需要通過循環才能找到合適的Entry一樣比較耗性能。但並不是說TreeMap性能低於HashMap就一無是處，TreeMap中的所有Entry總是按key根據指定的排序規則保持有序狀態。

備注:紅黑數是一種自平衡二叉查找樹 , 它們當中每一個節點的比較值都必須大於或等於在它的左子樹中的所有節點，並且小於或等於在它的右子樹中的所有節點。這確保紅黑樹運作時能夠快速的在樹中查找給定的值。

現在我們來觀察TreeMap的put(K key,V value)方法，該方法將Entry放入TreeMap的Entry鏈，並維護該Entry鏈的有序狀態。下面列出源碼:

public V put(K key, V value) {
      //定義一個t來保存根元素
        Entry<K,V> t = root;
        //如果t==null，表明是一個空鏈表
        if (t == null) {
        //如果根節點為null，將傳入的鍵值對構造成根節點（根節點沒有父節點，所以傳入的父節點為null）
            root = new Entry<K,V>(key, value, null);
            //設置該集合的size為1
            size = 1;
            //修改此時+1
            modCount++;
            return null;
        }
        // 記錄比較結果
        int cmp;
        Entry<K,V> parent;
        // 分割比較器和可比較接口的處理
        Comparator<? super K> cpr = comparator;
        // 有比較器的處理，即采用定制排序
        if (cpr != null) {
            // do while實現在root為根節點移動尋找傳入鍵值對需要插入的位置
            do {
                //使用parent上次循環後的t所引用的Entry
                // 記錄將要被摻入新的鍵值對將要節點(即新節點的父節點)
                parent = t;
                // 使用比較器比較父節點和插入鍵值對的key值的大小
                cmp = cpr.compare(key, t.key);
                // 插入的key較大
                if (cmp < 0)
                    t = t.left;
                // 插入的key較小
                else if (cmp > 0)
                    t = t.right;
                // key值相等，替換並返回t節點的value(put方法結束)
                else
                    return t.setValue(value);
            } while (t != null);
        }
        // 沒有比較器的處理
        else {
            // key為null拋出NullPointerException異常
            if (key == null)
                throw new NullPointerException();
            Comparable<? super K> k = (Comparable<? super K>) key;
            // 與if中的do while類似，只是比較的方式不同
            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);
        }
        // 沒有找到key相同的節點才會有下面的操作
        // 根據傳入的鍵值對和找到的“父節點”創建新節點
        Entry<K,V> e = new Entry<K,V>(key, value, parent);
        // 根據最後一次的判斷結果確認新節點是“父節點”的左孩子還是又孩子
        if (cmp < 0)
            parent.left = e;
        else
            parent.right = e;
        // 對加入新節點的樹進行調整
        fixAfterInsertion(e);
        // 記錄size和modCount
        size++;
        modCount++;
        // 因為是插入新節點，所以返回的是null
        return null;
    }

上面程序中的兩個do…while就是實現”排序二叉樹”的關鍵算法。每當程序希望添加新節點時，總是從樹的根節點開始比較，即將根節點當成當前節點。
如果新增節點大於當前節點且當前節點的右子節點存在，則以右子節點作為當前節點。並繼續循環

如果新增節點小於當前節點且當前節點的左子節點存在，則以左子節點作為當前節點。並繼續循環

如果新增節點等於當前節點，則新增節點覆蓋當前節點，並結束循環。

當TreeMap根據key來取出value時，TreeMap對應的方法如下:

public V get(Object key) {
     //根據key取出Entry
     Entry<K,V> p = getEntry(key);
     //取出Entry所包含的value
     return (p==null ? null : p.value);
 }

現在我們可以知道，其實get(Object key)方法實質上是由getEntry()方法實現的。現在我們來看getEntry(Object key)的源碼:

final Entry<K,V> getEntry(Object key) {
    // 如果有比較器，返回getEntryUsingComparator(Object key)的結果
    if (comparator != null)
        return getEntryUsingComparator(key);
    // 查找的key為null，拋出NullPointerException
    if (key == null)
        throw new NullPointerException();
    // 如果沒有比較器，而是實現了可比較接口
    //將key強制轉換為Comparable接口
    Comparable<? super K> k = (Comparable<? super K>) key;
    // 獲取根節點
    Entry<K,V> p = root;
    // 從根節點開始對樹進行遍歷查找節點
    while (p != null) {
        // 把key和當前節點的key進行比較
        int cmp = k.compareTo(p.key);
        // key小於當前節點的key
        if (cmp < 0)
            // p “移動”到左節點上
            p = p.left;
        // key大於當前節點的key
        else if (cmp > 0)
            // p “移動”到右節點上
p = p.right;
        // key值相等則當前節點就是要找的節點
        else
            // 返回找到的節點
            return p;
        }
    // 沒找到則返回null
    return null;
}

getEntry(Object obj)方法也是充分利用排序二叉樹的特性來搜索目標Entry。程序依然從二叉數的根節點開始，如果被搜索節點大於當前節點，程序向”右子樹”搜索，如果小於，則向”左子樹”搜索。如果相等則說明找到了指定節點。

我們觀察到當該TreeMap采用了定制排序。在采用定制排序的方式下，TreeMap采用getEntryUsingComparator(key)方法來根據key獲取Entry。

final Entry<K,V> getEntryUsingComparator(Object key) {
    K k = (K) key;
    // 獲取比較器
Comparator<? super K> cpr = comparator;
// 其實在調用此方法的get(Object key)中已經對比較器為null的情況進行判斷，這裡是防御性的判斷
if (cpr != null) {
    // 獲取根節點
        Entry<K,V> p = root;
        // 遍歷樹
        while (p != null) {
            // 獲取key和當前節點的key的比較結果
            int cmp = cpr.compare(k, p.key);
            // 查找的key值較小
            if (cmp < 0)
                // p“移動”到左孩子
                p = p.left;
            // 查找的key值較大
            else if (cmp > 0)
                // p“移動”到右節點
                p = p.right;
            // key值相等
            else
                // 返回找到的節點
                return p;
        }
}
// 沒找到key值對應的節點，返回null
    return null;
}

其實getEntry()和getEntryUsingComparator()這兩個方法實現思路幾乎完全類似。只是前者對自然排序的TreeMap獲取有效，後者對定制排序的TreeMap有效。

通過上述源碼其實不難看出，TreeMap這個工具類的實現其實很簡單。或者說，從本質上來說TreeMap就是一棵”紅黑樹”，每個Entry就是一個節點。