TreeMap 源碼解析

之前瞭解了ArrayList、LinkedList、HashMap、ConcurrentHashMap，接下來再來了解下TreeMap。TreeMap 底層直接維護了一個紅黑樹，按照慣例，先從繼承 及構造函數開始。

public class TreeMap<K,V> extends AbstractMap<K,V> implements NavigableMap<K,V>, Cloneable, java.io.Serializable
{
    // 比較器
    private final Comparator<? super K> comparator;

    // 跟節點
    private transient Entry<K,V> root;

    // 樹中實體（節點）的個數
    private transient int size = 0;

    // 對樹修改的次數
    private transient int modCount = 0;

    // 無參構造
    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);
    }

    // 有序Map構造
    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) {
        }
    }
}

構造函數及全局變量已經添加註釋，也是非常容易理解的，接下來以無參狗造爲基礎，通過添加元素，一步一步解析TreeMap。

一、put() 函數

先上源碼，通過源碼解釋更清楚明白（先屏蔽無用的代碼）：

    public V put(K key, V value) {
        Entry<K,V> t = root;
        if (t == null) {
            compare(key, key); // type (and possibly null) check
            // 第一次添加（跟節點爲 null），將新增的節點設置爲跟節點
            root = new Entry<>(key, value, null);
            size = 1;
            modCount++;
            return null;
        }
        ...
    }

從上述源碼可以清楚的看到，當初始化 TreeMap 後的第一次 put 操作時，root 爲 null ，直接把第一個 put 的數據設置爲跟節點。接下來繼續添加第二個元素，源碼如下：

    public V put(K key, V value) {
        Entry<K,V> t = root;
        if (t == null) {
           ...
        }
        int cmp;
        Entry<K,V> parent;
        // split comparator and comparable paths
        Comparator<? super K> cpr = comparator;
        if (cpr != null) {
           ...
        }
        else {
            if (key == null)
                throw new NullPointerException();
            @SuppressWarnings("unchecked")
                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 進行一個非空的判斷，說明在 TreeMap 中 key 不能爲 null 。在下面的 do while 中，從跟節點開始 依次和 本次添加的元素的 key 進行比較，從而找到存放該數據的節點位置。接下來，同樣根據默認的排序規則進行比較，將新創建的節點放到指定的位置上。由 TreeMap 的靜態內部類可知，新創建的節點 默認爲 黑色，源碼如下：

    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;

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

由紅黑樹的特性可知，每次新增加的節點均爲黑色是顯然不對的，所以在添加新的節點之後必須要對紅黑樹做一定的調整，使其成爲真正的紅黑樹，而調整紅黑樹的方法就是 fixAfterInsertion(e) 。可以看到它是由新添加的節點開始一步步調整的，源碼如下：

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

這個就不一步步梳理了，因爲整體來看還是比較簡單的，無非就是通過一系列的判斷進行左旋或者右旋使其符合紅黑樹的特性，想要了解紅黑樹 可以參考另一篇博文。

二、remove(K key); 方法

    public V remove(Object key) {
        Entry<K,V> p = getEntry(key);
        if (p == null)
            return null;

        V oldValue = p.value;
        deleteEntry(p);
        return oldValue;
    }

看起來也是非常簡單，通過key 獲取對應的節點，判斷節點是否爲空。而真正的刪除在 deleteEntry 中，源碼如下：

    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

        // 獲取刪除節點的左節點或者右節點
        Entry<K,V> replacement = (p.left != null ? p.left : p.right);
    
        if (replacement != 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;

            // 左右父節點都置爲空 以便fixAfterDeletion 使用.
            p.left = p.right = p.parent = null;

            // Fix replacement
            if (p.color == BLACK)
                fixAfterDeletion(replacement);
        } else if (p.parent == null) { // 刪除的節點 沒有子節點 也沒有父節點，說明刪除的是跟節點，直接將 root 設置爲 null 。
            root = null;
        } else { //  沒有子節點
            
            if (p.color == BLACK)
                fixAfterDeletion(p);
            // 如果有父節點，刪除節點爲左節點 或者 右節點 直接設置爲null
            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;
            }
        }
    }

刪除一個節點，首先判斷該節點有沒有左右節點，如果有，則獲取右節點的最左下角的那個節點並返回，源碼如下。

    /**
     * Returns the successor of the specified Entry, or null if no such.
     */
    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;
        }
    }

刪除節點之後 紅黑樹的轉變，源碼如下：

    private void fixAfterDeletion(Entry<K,V> 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);
    }

這個 刪除之後的轉變和添加新節點的轉變是非常類似的，可以參照一個紅黑樹進行解讀。TreeMap 的簡單瞭解就到這裏，更深層次的解讀，以後繼續分享。