从TreeMap源码搞懂红黑树

红黑树概述

在之前分析集合框架源码时笔者曾说过“红黑树的部分以后有经理再进行补充吧，左旋右旋，treefy，实在是磕不懂源码了哈哈哈”，终于狠下心从TreeMap源码仔细研究了一遍，来吧，好好分析分析红黑树这个牛逼的数据结构。

TreeMap源码最开始的注释中就写到了这样的一句话：

/**
 * A Red-Black tree based {@link NavigableMap} implementation.
 */

也就是说 TreeMap 是由红黑树实现的，红黑树和AVL树类似，都是在进行插入和删除元素时，通过旋转来保持自身平衡，红黑树只要求从根节点到叶子节点的最长路径不超过最短路径的2倍，即保持相对平衡即可（而AVL树要求左右子树高度差不超过1，姑且把这种平衡称作绝对平衡），红黑树在二叉搜索树的节点上增加了一个属性，即颜色，只可以是黑色或是红色，除了是二叉搜索树以外（二叉搜索树、AVL的概念就不在这里墨迹了吧…），它还有5个约束条件（“有红必有黑，红红不相连”）：

• 节点只能是黑色或红色；
• 根节点必须是黑色；
• 所有 NIL 节点都是黑色，NIL 即叶子节点下挂的两个虚节点；
• 一条路径上不允许出现相邻的红色节点；
• 在任何递归子树内，根节点到叶子节点的所有路径上包含数量相同的黑色节点。

红黑树的结构保证了新增、删除、查找的最坏时间复杂度均为 O(logn)，红黑树比AVL树的优势在于，AVL树在为了满足绝对平衡时，可能从最深层一直到第0层每层都不得不进行选择，而红黑树能够保证任何旋转在3次内完成。

红黑树的节点在 TreeMap 中（作为一个内部类）定义如下：

public class TreeMap<K,V> extends AbstractMap<K,V>
    implements NavigableMap<K,V>, Cloneable, java.io.Serializable {
    .......
        
	// Red-black mechanics
    private static final boolean RED   = false; // 红色节点
    private static final boolean BLACK = true;  // 黑色节点

    /**
	 * 红黑树中的节点
     * Node in the Tree.  Doubles as a means to pass key-value pairs back to
     * user (see Map.Entry).
     */

    static final class Entry<K,V> implements Map.Entry<K,V> {
        // fields
        K key;					// 键
        V value;				// 值
        Entry<K,V> left;		// 左孩子
        Entry<K,V> right;		// 右孩子
        Entry<K,V> parent;		// 父亲
        boolean color = BLACK;	// 节点颜色
        
        // methods
        ......
    }
    ......
}

这些红黑树节点按照 key 即键的大小组织成红黑树排列在 TreeMap 结构中。

TreeMap 依靠 Comparable / Comparator 来实现 key 的去重；

HashMap 依靠 hashCode 和 equals 实现去重。两者具有本质差别。

TreeMap中二叉搜索树的操作

先看 TreeMap 中是怎么查找节点的，TreeMap 中一共有6个（3组）查找节点 Entry 的方法。

• getEntry(Object key)：利用自然排序（extends Comparable）查找节点。
• getEntryUsingComparator(Object key)：利用定制排序（Comparator）查找节点。
• getCeilingEntry(K key)：
  • 存在给定键，返回给定键的节点；
  • 不存在给定键，返回比给定键大的最小值。
• getFloorEntry(K key)：
  • 存在给定键，返回给定键的节点；
  • 不存在给定键，返回比给定键小的最大值。
• getHigherEntry(K key)：下一个比给定键大的值。
• getLowerEntry(K key)：上一个比给定键小的值。

简单看下 getEntry(Object key) 的源码，就是典型的二叉搜索树的查找操作，因为红黑树本身就是二叉搜索树：

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); // 自然排序->compareTo
        if (cmp < 0)
            p = p.left;
        else if (cmp > 0)
            p = p.right;
        else
            return p;
    }
    return null;
}

而 getCeilingEntry(K key) 由于是要找比给定键大的最小值，故需要特殊处理操作：

final Entry<K,V> getCeilingEntry(K key) {
    Entry<K,V> p = root;
    while (p != null) {
        int cmp = compare(key, p.key);
		// 当所有元素都比给定键大时，则返回树中最左的节点
        if (cmp < 0) {
            if (p.left != null)
                p = p.left;
            else
                return p;
        } else if (cmp > 0) {
            if (p.right != null) {
                p = p.right;
            } else {
				/**
				 *     	 6 
				 *      / \
				 *	   2	
				 *	  	\	
				 *		 3
				 *		  \
				 *		   4
				 */
                Entry<K,V> parent = p.parent;
                Entry<K,V> ch = p;
				// 找比4大的最小值（应为6）
				// 结束时，ch == parent.left，即ch为2，parent为6
                while (parent != null && ch == parent.right) {
                    ch = parent;
                    parent = parent.parent;
                }
                return parent;
            }
        } else // 存在给定键，返回给定键的节点
            return p;
    }
    return null;
}

所以针对节点的查找，就是在一个二叉搜索树中如何查找节点。

TreeMap中平衡树的操作

平衡树的操作主要是左旋与右旋：

/**
 * 左旋（当前节点(22)下沉到左孩子处）
 *     
 *    /					 /
 *   22				    25
 *  /  \     左旋      /  \
 * 20  25    ---->    22  27
 *    /  \           /  \   \
 *   23  27         20  23  29 
 *        \
 *        29
 */
private void rotateLeft(Entry<K,V> p) {
	// 如果参数节点不为NIL
    if (p != null) { // p = 22
		// 当前节点的右子节点r
        Entry<K,V> r = p.right; // r = 25 
		// 将r的左子树设置为当前节点的右子树
        p.right = r.left; // 22.right = 25.left = 23
		// 若r的左子树不空，则将其父亲设置为当前节点
        if (r.left != null)
            r.left.parent = p; // 23.parent = 22
		// 将p的父亲设置为r的父亲
        r.parent = p.parent;
		
		// 无论如何，r都要在p父亲心中取代p的位置
        if (p.parent == null)
            root = r;
        else if (p.parent.left == p)
            p.parent.left = r;
        else // p.parent.right == p
            p.parent.right = r;
			
		// 将p设置为r的左子树	
        r.left = p; // 25.left = 22
		// 将r设置为p的父亲
        p.parent = r; // 22.parent = 25
    }
}

/**
 * 右旋（当前节点(17)下沉到右孩子处）
 *     
 * 		  /                 /
 *       17				   15
 *      /  \			  /  \
 *     15  18   右旋     13  17
 *    /  \	    ---->	/    / \
 *   13  16            11   16 18
 *  /
 * 11
 */
private void rotateRight(Entry<K,V> p) {
	// 如果参数节点不为NIL
    if (p != null) { // p = 17
		// 当前节点的左子节点l
        Entry<K,V> l = p.left; // l = 15
		// 将l的右子树设置为当前节点的左子树
        p.left = l.right; // 17.left = 15.right = 16
		// 若l的右子树不空，则将其父亲设置为当前节点
        if (l.right != null) l.right.parent = p; // 16.parent = 17
		// 将p的父亲设置为l的父亲
        l.parent = p.parent;
		
		// 无论如何，l都要在p父亲心中取代p的位置
        if (p.parent == null)
            root = l;
        else if (p.parent.right == p)
            p.parent.right = l;
        else p.parent.left = l;
		
		// 将p设置为l的右子树
        l.right = p; // 15.right = 17
		// 将l设置为p的父亲
        p.parent = l; // 17.right = 15
    }
}

TreeMap中红黑树的操作

private void fixAfterInsertion(Entry<K,V> x) {
	// 新节点一律先赋值为红色
    x.color = RED;

	// ① 新节点是根节点或其父节点为黑色，则插入红色节点不会破坏红黑树的性质，无需调整
	// ② 其父亲为红色节点，由于红黑树规定不能出现相邻的两个红色节点，故需重新着色或旋转
	// 退出循环时满足①，即x不断上游，直到父节点为黑色或已到达根节点
    while (x != null && x != root && x.parent.color == RED) {
		// 若父亲是爷爷的左孩子
        if (parentOf(x) == leftOf(parentOf(parentOf(x)))) {
			// 取决于右叔y
            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，进入到下一轮循环
                x = parentOf(parentOf(x));
			// 右叔为黑，需要旋转
            } else {
				// 若x是父亲的右孩子，则先对父亲做一次左旋
				// 转化为x是父亲左孩子的情况
                if (x == rightOf(parentOf(x))) {
					// 对父亲做左旋操作，父亲左沉到左侧位置取代为新的x
                    x = parentOf(x);
                    rotateLeft(x);
                }
				// 重新着色并对爷爷进行右旋
                setColor(parentOf(x), BLACK);
                setColor(parentOf(parentOf(x)), RED);
                rotateRight(parentOf(parentOf(x)));
            }
		// 父节点是爷爷的右孩子		
        } else {
			// 取决于左叔y
            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，进入到下一轮循环
                x = parentOf(parentOf(x));
			// 左叔为黑，需要旋转
            } else {
				// 若x是父亲的左孩子，则先对父亲做一次右旋
				// 转化为x是父亲右孩子的情况
                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;
}