TreeMap - 红黑树源码分析

结构

简单介绍

TreeMap本质上是一个红黑树，只不过节点存储的是一个key-value的键值对。
TreeMap继承自AbstractMap，同时实现了NavigableMap、Cloneable和Serializable接口。

源码剖析

基础结构

首先来看看TreeMap的基础结构

/**
     * The comparator used to maintain order in this tree map, or
     * null if it uses the natural ordering of its keys.
     *
     * @serial
     */
private final Comparator<? super K> comparator;

private transient Entry<K,V> root;

/**
     * The number of entries in the tree
     */
private transient int size = 0;

/**
     * The number of structural modifications to the tree.
     */
private transient int modCount = 0;

private transient EntrySet entrySet;
private transient KeySet<K> navigableKeySet;
private transient NavigableMap<K,V> descendingMap;

其中Comparator是比较器，用于比较Key的；root是指的这个红黑树的根，也就是说，这里是整个TreeMap真正存储值的地方；size是树的大小，modCount是对树的结构修改次数。这两个都是控制字段。
​

再来看看Entry的结构

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

可以看见，这里有个color的标记，说明我们的TreeMap里面的数据结构就是一棵红黑树，即利用红黑树原理实现的二叉搜索树。
​

建树过程

要学习红黑树，我们就来看看TreeMap的建树、插入和删除方法

private final Entry<K,V> buildFromSorted(int level, int lo, int hi,
                                             int redLevel,
                                             Iterator<?> it,
                                             java.io.ObjectInputStream str,
                                             V defaultVal)
        throws  java.io.IOException, ClassNotFoundException {

        if (hi < lo) return null;

        int mid = (lo + hi) >>> 1;

        Entry<K,V> left  = null;
        if (lo < mid)
            left = buildFromSorted(level+1, lo, mid - 1, redLevel,
                                   it, str, defaultVal);

        // extract key and/or value from iterator or stream
        K key;
        V value;
        if (it != null) {
            if (defaultVal==null) {
                Map.Entry<?,?> entry = (Map.Entry<?,?>)it.next();
                key = (K)entry.getKey();
                value = (V)entry.getValue();
            } else {
                key = (K)it.next();
                value = defaultVal;
            }
        } else { // use stream
            key = (K) str.readObject();
            value = (defaultVal != null ? defaultVal : (V) str.readObject());
        }

        Entry<K,V> middle =  new Entry<>(key, value, null);

        // color nodes in non-full bottommost level red
        if (level == redLevel)
            middle.color = RED;

        if (left != null) {
            middle.left = left;
            left.parent = middle;
        }

        if (mid < hi) {
            Entry<K,V> right = buildFromSorted(level+1, mid+1, hi, redLevel,
                                               it, str, defaultVal);
            middle.right = right;
            right.parent = middle;
        }

        return middle;
    }

其中，hi是当前这棵树的最大索引，lo是当前这棵树的最小索引，redLevel是节点应为红色的级别。
如果最大索引小于最小索引，这是递归建树的结束条件，说明到底了。
mid等于最大值+最小值除以2，说明这里使用了二分的方法建树。lo小于mid，在lo小于hi的情况下均成立。
lo<mid，先在左边建树。很明显，这是一个左中右建树过程。递归建立左边的树之后，开始处理中间的情况。
it是 m.entrySet().iterator()，即整个TreeMap的节点指针，如果该迭代器指针不为null，则需要从迭代器中取出当前节点的值，

 if (defaultVal==null) {
     Map.Entry<?,?> entry = (Map.Entry<?,?>)it.next();
     key = (K)entry.getKey();
     value = (V)entry.getValue();
 } else {
     key = (K)it.next();
     value = defaultVal;
 }

这部分就是在取迭代器中的值。可以发现这里的key可以不是字符串。
如果迭代器为null，那么就从字符流里面获取key-value。
然后创建一个中间节点，即自己：

Entry<K,V> middle =  new Entry<>(key, value, null);

然后给自己标记颜色

if (level == redLevel)
    middle.color = RED;

这里的规则就是如果当前的级别等于红色的级别，那么就给节点标记为红色。即，在递归传入参数的时候就确定了哪一级别就是红色级别，其余都是黑色级别。
然后，再把左右子树给组装到根节点上。

if (left != null) {
    middle.left = left;
    left.parent = middle;
}

if (mid < hi) {
    Entry<K,V> right = buildFromSorted(level+1, mid+1, hi, redLevel,
                                       it, str, defaultVal);
    middle.right = right;
    right.parent = middle;
}

这下，整个建树函数就分析完了，现在我们需要来看看一些细节，首先是TreeMap是怎么确定的每棵树的redLevel的？
在TreeMap中，使用的是 _computeRedLevel _函数来计算应该为红色节点的深度。

/**
     * Find the level down to which to assign all nodes BLACK.  This is the
     * last `full' level of the complete binary tree produced by
     * buildTree. The remaining nodes are colored RED. (This makes a `nice'
     * set of color assignments wrt future insertions.) This level number is
     * computed by finding the number of splits needed to reach the zeroeth
     * node.  (The answer is ~lg(N), but in any case must be computed by same
     * quick O(lg(N)) loop.)
     */
private static int computeRedLevel(int sz) {
    int level = 0;
    for (int m = sz - 1; m >= 0; m = m / 2 - 1)
        level++;
    return level;
}

让我们来理解一下官方的注释：本质上就是获取了这棵树如果是一棵完全二叉树，那么最树高为多少。
等于 size+1的底数为2的对数值的向下取整的值。
我们来看看，整个建树的调用过程

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) {
    }
}

private void buildFromSorted(int size, Iterator<?> it,
                             java.io.ObjectInputStream str,
                             V defaultVal)
    throws  java.io.IOException, ClassNotFoundException {
    this.size = size;
    root = buildFromSorted(0, 0, size-1, computeRedLevel(size),
                           it, str, defaultVal);
}

也就是说，在TreeMap初次建立的时候，是将所有值放到EntrySet里面或者实在一个字符流里面，然后利用中序建树的方式建树，将level等于树高的层置为红色节点，其余全是黑色节点。

用思维导图看这个过程就是：

不过，值得注意的是，这只针对与TreeMap以SortedMap为基础建树的过程，如果，本身不是一个排序过后的Map，TreeMap还是只能以有一个基础树的情况下，插入建树。不过这个过程涉及向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

        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();
        @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;
}

如果根节点为null，则把现在要put的值给根节点，否则就通过比较器不断的寻找到一个为null的地方

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

这个部分就是判断当前插入到TreeMap里面的值有没有默认实现比较器函数，如果有，就使用值里面的，这也是为什么我们在使用TreeMap的时候，放进去的值要求是可以被比较的。如果没有，就使用默认的比较器。
在找到了该值存储的位置之后，就可以开始着手将值放到该位置了。

Entry<K,V> e = new Entry<>(key, value, parent);
if (cmp < 0)
    parent.left = e;
else
    parent.right = e;

然后，我们需要去判断整棵树是否是符合红黑树规则的，所以这里调用了fixAfterInsertion函数。

上色旋转过程

fixAfterInsertion函数实现了上色过程，其中下面的函数，实现了红黑树的旋转过程。

其中左旋和右旋函数代码为：

/** From CLR */
private void rotateLeft(Entry<K,V> p) {
    if (p != null) {
        Entry<K,V> r = p.right;
        p.right = r.left;
        if (r.left != null)
            r.left.parent = p;
        r.parent = p.parent;
        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;
    }
}

/** From CLR */
private void rotateRight(Entry<K,V> p) {
    if (p != null) {
        Entry<K,V> l = p.left;
        p.left = l.right;
        if (l.right != null) l.right.parent = p;
        l.parent = p.parent;
        if (p.parent == null)
            root = l;
        else if (p.parent.right == p)
            p.parent.right = l;
        else p.parent.left = l;
        l.right = p;
        p.parent = l;
    }
}

其中左旋部分的逻辑如下（按代码执行步骤分析）：
r = p.right

p.right = r.left;
if (r.left != null)
r.left.parent = p;

​

此时，r和p原本的一条关系就断掉了。
r.parent = p.parent;

if (p.parent == null)
root = r;
else if (p.parent.left == p)
p.parent.left = r;
else
p.parent.right = r;
这一部分是判断p是父亲的左右哪个子节点，然后替换成r。

r.left = p;
p.parent = r;

左旋函数的最终转换结果的图示如下：
​

右旋函数的最终转换结果的图示如下：

它这个左旋右旋，效果似乎和平衡二叉树不太一样。看起来有点像左右变右左，右左变左右。不过这正也是红黑树比平衡二叉树快的地方，即它不需要去考虑底部节点的平衡性。那么我们仔细看看调用它的方法究竟为啥要这么进行旋转吧。
​

上色过程：

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

图解上述过程：
首先

if (colorOf(y) == RED) {
    setColor(parentOf(x), BLACK);
    setColor(y, BLACK);
    setColor(parentOf(parentOf(x)), RED);
    x = parentOf(parentOf(x));
} else {
    ... ...
}

这个条件的转换如下图：
​

此时，将x = parentOf(parentOf(x));，因为从图中我们可以发现，最顶部下面的节点已经符合红黑树的性质了，但是由于改变了根节点的性质，那么这个根节点的上层节点之间是否存在问题，就需要进行回溯。所以令x = parentOf(parentOf(x));。
再来看看

if (colorOf(y) == RED) {
    ... ...
} else {
    if (x == rightOf(parentOf(x))) {
        x = parentOf(x);
        rotateLeft(x);
    }
    setColor(parentOf(x), BLACK);
    setColor(parentOf(parentOf(x)), RED);
    rotateRight(parentOf(parentOf(x)));
}

这部分内容的转换如下图：

然后再对跟节点进行右旋，示意图如下：

而当我们去掉中间过程，就是如下图：

​

​

而最新的x=parentOf(x)。
​

然后再来看看else里面的内容

} 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 {
        ... ...
    }
}

图示如下（设pp(x) 为 parentOf(parentOf(x))，p(x) 为 parentOf(x)）:

然后令 x = ppx，进入回溯。
​

最后看看剩下的那部分

} else {
    Entry<K,V> y = leftOf(parentOf(parentOf(x)));
    if (colorOf(y) == RED) {
       ... ...
    } else {
        if (x == leftOf(parentOf(x))) {
            x = parentOf(x);
            rotateRight(x);
        }
        setColor(parentOf(x), BLACK);
        setColor(parentOf(parentOf(x)), RED);
        rotateLeft(parentOf(parentOf(x)));
    }
}

​

这四个过程被一个while循环包裹着，就是利用回溯的方法保证整棵树瞒住红黑树的性质。

while (x != null && x != root && x.parent.color == RED) {
	... ...
}

删除节点

源码是这样写的

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

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

其中，getEntry方法就是查找方法，代码如下：

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();
    @SuppressWarnings("unchecked")
    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;
}

这就是二叉查找树的查找方法了，递归查找，时间复杂度为log(n+1)，即树的高度。
deleteEntry函数，会改变树结构，所以需要重新选择和着色，源代码如下：

/**
     * Delete node p, and then rebalance the tree.
     */
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;
        }
    }
}

由于要干掉一个节点，所以需要对其左右子树进行处理。如果该节点左右子树都存在，则调用successor(p)找到谁来占用P的位置。根据该函数的实现，找的结果，如果它所在的节点存在右子树，那么找到的就是它的右子树的最小值，如果它所在的节点不存在右子树，那么找到的就是向右单调向上的最后一个节点。
​

替换完了之后，这棵红黑树有可能是不符合部分性质的，所以需要调用下面的修复函数进行修改。这与修复插入类似。

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