寫本文的原由是昨晚做夢居然夢到了在看原始碼，於是便有了此文......

雖然標題是關於紅黑樹的，不過此文是結合圖片，通過分析TreeMap的原始碼，讓你理解起來不是那麼枯燥（前方高能，此文圖片眾多，慎入）。

作者 | 馬雲飛

責編 | 胡巍巍

概述

TreeMap是紅黑樹的Java實現，紅黑樹能保證增、刪、查等基本操作的時間複雜度為O(lgn)。    

首先我們來看一張TreeMap的繼承體系圖：

image

還是比較直觀的，這裡來簡單說一下繼承體系中不常見的介面NavigableMap和SortedMap，這兩個介面見名知意。先說NavigableMap介面，NavigableMap介面聲明瞭一些列具有導航功能的方法，比如：

/**
 * 返回紅黑樹中最小鍵所對應的 Entry
 */
Map.Entry<K,V> firstEntry();

/**
 * 返回最大的鍵 maxKey，且 maxKey 僅小於引數 key
 */
K lowerKey(K key);

/**
 * 返回最小的鍵 minKey，且 minKey 僅大於引數 key
 */
K higherKey(K key);

// 其他略

通過這些導航方法，我們可以快速定位到目標的Key或Entry。至於 SortedMap介面，這個介面提供了一些基於有序鍵的操作，比如：

/**
 * 返回包含鍵值在 [minKey, toKey) 範圍內的 Map
 */
SortedMap<K,V> headMap(K toKey);();

/**
 * 返回包含鍵值在 [fromKey, toKey) 範圍內的 Map
 */
SortedMap<K,V> subMap

(K fromKey, K toKey);

// 其他略

以上就是兩個介面的介紹，很簡單。關於TreeMap的繼承體系就這裡就說到這，接下來我們深入原始碼進行分析。

原始碼分析

新增

紅黑樹最複雜的無非就是增刪了，這邊我們先介紹增加一個元素，瞭解紅黑樹的都知道，當往 TreeMap 中放入新的鍵值對後，可能會破壞紅黑樹的性質。首先我們先鞏固一下紅黑樹的特性。

• 節點是紅色或黑色。

• 根節點是黑色。

• 每個葉子節點都是黑色的空節點（NIL節點）。

• 每個紅色節點的兩個子節點都是黑色（從每個葉子到根的所有路徑上不能有兩個連續的紅色節點）。

• 從任一節點到其每個葉子的所有路徑都包含相同數目的黑色節點。

接下來我們看看新增到底做了什麼處理：

  public V put(K key, V value) {
        TreeMapEntry<K,V> t = root;
        if (t == null) {

            if (comparator != null) {
                if (key == null) {
                    comparator.compare(key, key);
                }
            } else {
                if (key == null) {
                    throw new NullPointerException("key == null");
                } else if (!(key instanceof Comparable)) {
                    throw new ClassCastException(
                            "Cannot cast" + key.getClass().getName() + " to Comparable.");
                }
            }
            root = new TreeMapEntry<>(key, value, null);
            size = 1;
            modCount++;
            return null;
        }
        int cmp;
        TreeMapEntry<K,V> parent;
        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 {
                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);
        }
        TreeMapEntry<K,V> e = new TreeMapEntry<>(key, value, parent);
        if (cmp < 0)
            parent.left = e;
        else
            parent.right = e;
        fixAfterInsertion(e);
        size++;
        modCount++;
        return null;
    }

這邊會先把根節點暫存依賴，如果根節點為null，則講新增的這個節點設為根節點。

如果初始化的時候指定了Comparator比較器，則講其插入到指定位置，否則使用key進行比較並且插入。

不斷地進行比較，找到沒有子節點的節點，將其插入到相應節點（注：如果查找出有相同的值只會更新當前值，CMP小於0是沒有左節點，反之沒有右節點）。

新插入的樹破環的紅黑樹規則，我們會通過fixAfterInsertion去進行相應的調整，這也是TreeMap插入實現的重點，具體我們看看他是怎麼通過Java實現的。

 private void fixAfterInsertion(TreeMapEntry<K,V> x) {
        x.color = RED;

        while (x != null && x != root && x.parent.color == RED) {
            if (parentOf(x) == leftOf(parentOf(parentOf(x)))) {
                TreeMapEntry<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 {
                TreeMapEntry<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;
    }

首先將新插入的節點設定為紅色，這邊做了一個判斷，新節點不為null，新節點不為根節點並且新節點的父節點為紅色。才會進入內部的判斷，因為其本身就是一個紅黑樹。

如果新節點的父節點為黑色，則他依舊滿足紅黑樹的特性。所以當其父節點為紅色進入內部的判斷。

如果新節點是其祖父節點的左子孫，則拿到其祖父節點的右兒子，也就是新節點的叔叔。如果叔叔節點是紅色。

則將其父節點設為黑色，講叔父節點設為黑色，然後講新節點直接其祖父節點。

否則如果新節點是其父節點的右節點，以其父節點進行左轉，將父節點設為黑色，祖父節點設為紅色，在通過祖父節點進行右轉。

else內容和上述基本一致。可以自己分析。最後我們還需要將跟節點設為黑色。

我們稍微看一下，他是怎麼進行左轉和右轉的。

// 右旋與左旋思路一致，只分析其一
// 結果相當於把p和p的兒子調換了
private void rotateLeft(Entry<K,V> p) {
    if (p != null) {
        // 取出p的右兒子
        Entry<K,V> r = p.right;
        // 然後將p的右兒子的左兒子，也就是p的左孫子變成p的右兒子
        p.right = r.left;
        if (r.left != null)
            // p的左孫子的父親現在是p
            r.left.parent = p;

        // 然後把p的父親，設定為p右兒子的父親
        r.parent = p.parent;
        // 這說明p原來是root節點
        if (p.parent == null)
            root = r;
        else if (p.parent.left == p)
            p.parent.left = r;
        else
            p.parent.right = r;
        r.left = p;
        p.parent = r;
    }
}

//和左旋類似
private void rotateRight(Entry<K,V> p) {
    // ...
}

下面我們通過圖解來看看如何插入一顆紅黑樹。

現有陣列int[] a = {1, 10, 9, 2, 3, 8, 7, 4, 5, 6};我們要將其變為一棵紅黑樹。

首先插入1，此時樹是空的，1就是根結點，根結點是黑色的：

然後插入元素10，此時依然符合規則，結果如下：

當插入元素9時，這時是需要調整的第一種情況，結果如下：

紅黑樹規則4中強調不能有兩個相鄰的紅色結點，所以此時我們需要對其進行調整。

調整的原則有多個相關因素，這裡的情況是，父結點10是其祖父結點1（父結點的父結點）的右孩子，當前結點9是其父結點10的左孩子，且沒有叔叔結點（父結點的兄弟結點），此時需要進行兩次旋轉，第一次，以父結點10右旋：

然後將父結點（此時是9）染為黑色，祖父結點1染為紅色，如下所示：

然後以祖父結點1左旋：

下一步，插入元素2，結果如下：

此時情況與上一步類似，區別在於父結點1是祖父結點9的左孩子，當前結點2是父結點的右孩子，且叔叔結點10是紅色的。

這時需要先將叔叔結點10染為黑色，再進行下一步操作，具體做法是將父結點1和叔叔結點10染為黑色，祖父結點9染為紅色，如下所示：

由於結點9是根節點，必須為黑色，將它染為黑色即可：

下一步，插入元素3，如下所示：

這和我們之前插入元素10的情況一模一樣，需要將父結點2染為黑色，祖父結點1染為紅色，如下所示：

然後左旋：

下一步，插入元素8，結果如下：

此時和插入元素2有些類似，區別在於父結點3是右孩子，當前結點8也是右孩子，這時也需要先將叔叔結點1染為黑色，具體操作是先將1和3染為黑色，再將祖父結點2染為紅色，如下所示：

此時樹已經平衡了，不需要再進行其他操作了，現在插入元素7，如下所示：

這時和之前插入元素9時一模一樣了，先將7和8右旋，如下所示：

然後將7染為黑色，3染為紅色，再進行左旋，結果如下：

下一步要插入的元素是4，結果如下：

這裡和插入元素2是類似的，先將3和8染為黑色，7染為紅色，如下所示：

但此時2和7相鄰且顏色均為紅色，我們需要對它們繼續進行調整。這時情況變為了父結點2為紅色，叔叔結點10為黑色，且2為左孩子，7為右孩子，這時需要以2左旋。

這時左旋與之前不同的地方在於結點7旋轉完成後將有三個孩子，結果類似於下圖：

這種情況處理起來也很簡單，只需要把7原來的左孩子3，變成2的右孩子即可，結果如下：

然後再把2的父結點7染為黑色，祖父結點9染為紅色。結果如下所示：

此時又需要右旋了，我們要以9右旋，右旋完成後7又有三個孩子，這種情況和上述是對稱的，我們把7原有的右孩子8，變成9的左孩子即可，如下所示：

下一個要插入的元素是5，插入後如下所示：

有了上述一些操作，處理5變得十分簡單，將3染為紅色，4染為黑色，然後左旋，結果如下所示：

最後插入元素6，如下所示：

又是叔叔結點3為紅色的情況，這種情況我們處理過多次了，首先將3和5染為黑色，4染為紅色，結果如下：

此時問題向上傳遞到了元素4，我們看2、4、7、9的顏色和位置關係，這種情況我們也處理過，先將2和9染為黑色，7染為紅色，結果如下：

最後7是根結點，染為黑色即可，最終結果如下所示：

可以看到，在插入元素時，叔叔結點是主要影響因素，待插入結點與父結點的關係決定了是否需要多次旋轉。

刪除

除了新增操作，紅黑樹的刪除也是很麻煩的…...我們看看怎麼通過Java去實現紅黑樹的刪除。具體程式碼如下：

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

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

其內部是通過Delete Entry去進行刪除的。所以我們具體看看Delete Entry的實現。

 private void deleteEntry(TreeMapEntry<K,V> p) {
        modCount++;
        size--;

        if (p.left != null && p.right != null) {
            TreeMapEntry<K,V> s = successor(p);
            p.key = s.key;
            p.value = s.value;
            p = s;
        } 

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

            p.left = p.right = p.parent = null;

            // Fix replacement
            if (p.color == BLACK)
                fixAfterDeletion(replacement);
        } else if (p.parent == null) { 
            root = null;
        } else {
            if (p.color == BLACK)
                fixAfterDeletion(p);

            if (p.parent != null) {
                if (p == p.parent.left)
  &nbs