前言

本文所有代码基于 JDK 1.8

主要分析最基本的构造方法及增删改查操作，最后在此基础上分析一些其他的方法。

成员变量

// 为 null 表示自然排序（从小到大）
private final Comparator<? super K> comparator;

// 红黑树（Entry 是红黑树节点）的根节点
private transient Entry<K,V> root;

// 元素个数
private transient int size = 0;

// 结构性修改次数
private transient int modCount = 0;

其中，Comparator 是一个接口，可用于自定义排序方式，接口中有一个 compare 方法，该方法返回一个 int 值，返回值为正数，零，负数各代表大于，等于，小于。

Entry 用于存储元素，下面看下它的实现：

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;

      Entry(K key, V value, Entry<K,V> parent) {
          this.key = key;
          this.value = value;
          this.parent = parent;
      }

      public K getKey() {
          return key;
      }

      public V getValue() {
          return value;
      }

      public V setValue(V value) {
          V oldValue = this.value;
          this.value = value;
          return oldValue;
      }

      public boolean equals(Object o) {
          if (!(o instanceof Map.Entry))
              return false;
          Map.Entry<?,?> e = (Map.Entry<?,?>)o;

          return valEquals(key,e.getKey()) && valEquals(value,e.getValue());
      }

      public int hashCode() {
          int keyHash = (key==null ? 0 : key.hashCode());
          int valueHash = (value==null ? 0 : value.hashCode());
          return keyHash ^ valueHash;
      }

      public String toString() {
          return key + "=" + value;
      }
  }

代码还是很好懂的，显然，Entry 是一个红黑树节点，持有左右孩子及父节点的引用。

了解了基本的成员变量，接下来看它的几个构造方法：

构造方法

TreeMap()

1
2
3

public TreeMap() {
    comparator = null;
}

默认构造方法，将 comparator 设置为 null，元素的排序方式为自然排序。

TreeMap(Comparator<? super K> comparator)

1
2
3

public TreeMap(Comparator<? super K> comparator) {
    this.comparator = comparator;
}

带 Comparator 参数，可自定义排序。

TreeMap(Map<? extends K, ? extends V> m)

public TreeMap(Map<? extends K, ? extends V> m) {
    comparator = null;
    putAll(m);
}

该方法将另一个 Map 的所有元素添加到当前 TreeMap

TreeMap(SortedMap<K, ? extends V> m)

public TreeMap(SortedMap<K, ? extends V> m) {
    comparator = m.comparator();
    try {
        buildFromSorted(m.size(), m.entrySet().iterator(), null, null);
    } 
    // ...
}

和上一个构造方法类似，该方法传入一个 SortedMap（TreeMap 就是其中一个 SortedMap），要求该 SortedMap 的 key 类型要和自己的 key 类型一致。

简单看过各个构造方法，接下来就分析一些基本的增删改查操作，先看 put 方法：

put

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

// 创建根节点
      if (t == null) {
	// 如果不指定 Comparator，或指定的 Comparator 没有处理空指针异常
	// 则 key 不能为 null，否则在比较元素时会引发空指针异常
          compare(key, key);
	
          root = new Entry<>(key, value, null);
          size = 1;
          modCount++;
          return null;
      }

      int cmp;
      Entry<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();

	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;

// （1）
// 调整成规范的红黑树
      fixAfterInsertion(e);

      size++;
      modCount++;
      return null;
  }

要注意，如果不指定 Comparator，或指定的 Comparator 没有处理空指针异常，则 key 不能为 null，否则在比较元素时会引发空指针异常。

总结下添加元素的过程：

如果还没有节点就创建根节点，并返回 null。
找到插入节点的父节点和插入位置，如果在查找过程出现比较结果相等的节点，更新该节点的值并返回旧的值。
构造节点并插入到相应位置。
调整成规范的红黑树。

下面分析下第四步，也就是注释（1）处的 fixAfterInsertion 方法：

fixAfterInsertion

由于红黑树调整过程涉及到左旋和右旋，它们的实现分别对应 rotateLeft 和 rotateRight 方法，看下 rotateLeft（rotateRight 就不贴了，和 rotateLeft 类似，只是左右反过来而已）

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

代码我就不用文字说明了，免得把人绕晕，直接上图，这样更加直观：

右旋原理也一样，如图：

了解了左旋和右旋，接下来看 fixAfterInsertion 的代码实现：

 private void fixAfterInsertion(Entry<K,V> x) {
   // 将插入节点置为红色节点
      x.color = RED;

      while (x != null && x != root && x.parent.color == RED) {
	// 如果 x 的父节点（设为 xp）是其父节点（设为 xpp）的左孩子
          if (parentOf(x) == leftOf(parentOf(parentOf(x)))) {
		// y 为 xpp 的右孩子，即 xp 的右兄弟
              Entry<K,V> y = rightOf(parentOf(parentOf(x)));
		// y 为红色节点，就将 xp，y 置为黑色，将 xpp 置为红色，最后 x 置为 xpp
              if (colorOf(y) == RED) {
                  setColor(parentOf(x), BLACK);
                  setColor(y, BLACK);
                  setColor(parentOf(parentOf(x)), RED);
                  x = parentOf(parentOf(x));
              } else {	// y 为 黑色节点或 y 为 null
			// 如果 x 是其父节点的右孩子，先将 x 置为 xp，然后对 x 进行左旋操作
                  if (x == rightOf(parentOf(x))) {
                      x = parentOf(x);
                      rotateLeft(x);
                  }
			// 将 xp 置为黑色，xpp 置为红色，以 xpp 为中心进行右旋操作
                  setColor(parentOf(x), BLACK);
                  setColor(parentOf(parentOf(x)), RED);
                  rotateRight(parentOf(parentOf(x)));
              }
          } 
	// 如果 x 的父节点（设为 xp）是其父节点（设为 xpp）的右孩子
	else {
		// y 为 xpp 的左孩子，即 xp 的左兄弟
              Entry<K,V> y = leftOf(parentOf(parentOf(x)));
		// y 为红色节点，就将 xp，y 置为黑色，将 xpp 置为红色，最后 x 置为 xpp
              if (colorOf(y) == RED) {
                  setColor(parentOf(x), BLACK);
                  setColor(y, BLACK);
                  setColor(parentOf(parentOf(x)), RED);
                  x = parentOf(parentOf(x));
              } else {    // y 为 黑色节点或 y 为 null
			// 如果 x 是其父节点的左孩子，先将 x 置为 xp，然后对 x 进行右旋操作
                  if (x == leftOf(parentOf(x))) {
                      x = parentOf(x);
                      rotateRight(x);
                  }
			// 将 xp 置为黑色，xpp 置为红色，以 xpp 为中心进行左旋操作
                  setColor(parentOf(x), BLACK);
                  setColor(parentOf(parentOf(x)), RED);
                  rotateLeft(parentOf(parentOf(x)));
              }
          }
      }
// 根节点置为黑色节点
      root.color = BLACK;
  }

过程的话注释都有说明，不过这样并不直观，还是用图来说话。

假设依次插入 1, 3, 5, 7, 9, 11, 13, 15。这是一个升序序列，如果不进行调整的话，节点都会堆在右子树，查找效率就变为 O(n) 了。

先看下依次插入 1, 3, 5, 7, 9, 11, 13 后形成的红黑树：

接下来看下插入节点 15 的过程：

在插入节点 15 的时候，需要在 while 循环中循环两次，也就是有两步，如果不懂得为什么要这样调整。

如果看不懂图示的过程，可以自己跟着代码走一遍。

看完红黑树的调整过程，回到 TreeMap，刚才看完了 put 方法，接下来就看下 remove 方法：

remove

  public V remove(Object key) {
// （1）
// 获得删除节点，没有该节点则返回 null
      Entry<K,V> p = getEntry(key);
      if (p == null)
          return null;

      V oldValue = p.value;

// （2）
// 删除相应节点
      deleteEntry(p);

      return oldValue;
  }

删除节点的第一步是获得删除节点，具体看注释（1）处的 getEntry 方法：

getEntry

  final Entry<K,V> getEntry(Object key) {
// 设置了 Comparator 时，获取过程也类似，只不过在比较时利用的是 Comparator 
      if (comparator != null)
          return getEntryUsingComparator(key);

// 查找时 key 不能为 null
// （其实如果没有设置 Comparator 的话，插入的时候已经不能为 null 了）
      if (key == null)
          throw new NullPointerException();
      
// 没有设置 Comparator 时根据自然排序（Comparable）的结果
// 在红黑树中找到相应的节点，若没有找到则返回 null
@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;
  }

根据 key 来获得节点，利用 Comparator 或 Comparable 比较 key 后，如果返回结果为 0，说明 key 相等，该 key 对应节点即为所找节点。

回到 remove 方法，继续看注释（2）处的 deleteEntry 方法，该方法删除相应节点：

deleteEntry

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

// 如果删除节点左右孩子皆不为空，先给它找到一个替代节点 s
// s 为 p 的右子树的深度最大的左孩子（即在平衡二叉树中，删除节点的下一节点）
// 然后将被删除节点的 key 和 value 置为 s 的 key 和 value，并将 p 置为 s
      if (p.left != null && p.right != null) {
          Entry<K,V> s = successor(p);
          p.key = s.key;
          p.value = s.value;
          p = s;
      }

// 若 p 的左孩子不为空，则 replacement 为 p 的左孩子，否则为 p 的右孩子
      // 注意：执行到这里时，p 最多只能有一个孩子（没有孩子也可能）  
      Entry<K,V> replacement = (p.left != null ? p.left : p.right);
      
// 节点 p 有孩子
      if (replacement != null) {
          // 节点 p 的孩子顶替节点 p 的位置
          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
          p.left = p.right = p.parent = null;

          // 如果节点 p 为黑色，需要调整删除了节点 p 后的树
          if (p.color == BLACK)
              fixAfterDeletion(replacement);
      } 
// 只有一个节点时，将根节点置为空
else if (p.parent == null) { 
          root = null;
      } 
// 节点 p 没有孩子
else { 
          // 如果节点 p 为黑色，调整红黑树
          if (p.color == BLACK)
              fixAfterDeletion(p);
	// 删除节点 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;
          }
      }
  }

这个方法比较长，重点是读懂它的逻辑。总的来说，该方法会根据删除节点所处的位置分了几种情况：

删除节点的左右孩子皆不为空
删除节点只有左孩子或只有右孩子
删除节点没有孩子

不同情况有不同的操作，但最终都是由两个操作组成：删除对应节点，以及调整红黑树。

下面还是用图来说明，以上面生成的红黑树为例：

首先分析删除节点的左右孩子皆不为空的情况，这里删除节点 11：

最后还要调整红黑树，调整操作在 fixAfterDeletion 方法中，之后再分析这个方法，这里先看删除部分。

还有另外两种情况，删除节点只有左孩子或只有右孩子，以及删除节点没有孩子，这两种情况的删除比较简单，分别以删除节点 13 和节点 9 为例：

在删除节点 13 时，最后还要进行调整，这里放在后面分析。在删除节点 9 时，需要在删除操作前先进行调整工作，这里直接给出了调整后的结果。

现在就来看一下删除后红黑树的调整，该操作在 fixAfterDeletion 方法中：

fixAfterDeletion

  private void fixAfterDeletion(Entry<K,V> x) {
// 如果 x 不是根节点，并且为黑色节点
      while (x != root && colorOf(x) == BLACK) {
	// x 是其父节点的左孩子
          if (x == leftOf(parentOf(x))) {
              Entry<K,V> sib = rightOf(parentOf(x));
		// x 的右兄弟为红色节点
              if (colorOf(sib) == RED) {
                  setColor(sib, BLACK);
                  setColor(parentOf(x), RED);
                  rotateLeft(parentOf(x));
                  sib = rightOf(parentOf(x));
              }
		// 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;
              }
          } 
          // x 是其父节点的右孩子，和上面操作相反
          else {
              // ...
          }

      }

// 将节点 x 置为黑色
      setColor(x, BLACK);
  }

这里的调整操作比较复杂，直接以上面的几个删除例子说明，在删除节点 11 和节点 13 的时候比较简单，传入的节点都是节点 15，该节点为红色，所以不用进入 while 循环中，直接将节点置为黑色即可。

现在看一下删除节点 9 前的调整，这里传入的节点是要删除的节点，即节点 9：

调整后从红黑树中删除节点 9 即可。

看完添加和删除元素，增删改查还剩下改和查，分别对应 replace 和 get 方法，这个方法比较简单，下面分别看下：

replace

replace 有两个重载方法，分别传入两个参数和三个参数，先看下传入两个参数的：

@Override
public V replace(K key, V value) {
    Entry<K,V> p = getEntry(key);
    if (p!=null) {
        V oldValue = p.value;
        p.value = value;
        return oldValue;
    }
    return null;
}

该方法传入要更改元素的 key 和新的 value，返回旧的 value，其中在查找节点的时候用到了 getEntry 方法，该方法在分析删除节点的时候看过了，这里不再多说。

另一个带有三个参数的方法如下：

@Override
public boolean replace(K key, V oldValue, V newValue) {
    Entry<K,V> p = getEntry(key);
    if (p!=null && Objects.equals(oldValue, p.value)) {
        p.value = newValue;
        return true;
    }
    return false;
}

该方法多了个参数表示旧 value，根据 key 和 value 来更改元素，更改成功就返回 true。

最后看下 get 方法，该方法就更简单了：

get

public V get(Object key) {
    Entry<K,V> p = getEntry(key);
    return (p==null ? null : p.value);
}

同样是通过 getEntry 方法找到相应节点，返回该节点的 value，没有找到则返回 null。

其他

下面看一下除了增删改查外的其他方法。

LeetCode 220 题的官方题解上使用到了 TreeSet 的 ceiling 和 floor 方法，而 TreeSet 是基于 TreeMap 实现的，所以这两个方法分别调用了 TreeMap 的 ceilingKey 和 floorKey 方法，所以在这里特别分析一下这两个方法：

ceilingKey

1
2
3

public K ceilingKey(K key) {
    return keyOrNull(getCeilingEntry(key));
}

getCeilingEntry 方法如下：

 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
		// 找不到相应的 key，返回传入 key 的下一节点
                 return p;
         } else if (cmp > 0) {
             if (p.right != null) {
                 p = p.right;
             } else {
		// 找不到相应的 key 时，寻找传入 key 的下一节点
		// 这里是寻找当前 key 所在左子树的最小父节点，不存在则返回 null
                 Entry<K,V> parent = p.parent;
                 Entry<K,V> ch = p;
                 while (parent != null && ch == parent.right) {
                     ch = parent;
                     parent = parent.parent;
                 }
                 return parent;
             }
         } 
// 找到相等的 key，返回该节点
else
             return p;
     }
     return null;
 }

该方法根据传入的 key 寻找相应节点，如果没有找到相应节点，就返回比该 key 稍大的下一节点，也没有就返回 null。

所以 ceilingKey 方法就是传入一个 key，如果 TreeMap 存在该 key，则返回该 key，不存在则返回比该 key 稍大的下一个 key。如果传入的 key 是最大的 key，则返回 null。

而 TreeSet 的 ceiling 方法底层就是调用了 TreeMap 的 ceilingKey 方法，只不过把 key 换成了 TreeSet 的元素而已。所以 TreeSet 的 ceiling 方法就是返回传入元素或者是传入元素的下一元素，不存在比传入元素更大的元素则返回 null。

floorKey

1
2
3

public K floorKey(K key) {
    return keyOrNull(getFloorEntry(key));
}

其中 getFloorEntry 和 getCeilingEntry 方法类似，只是在查找不到相应 key 时，返回比该 key 较小的上一节点，这里就不再分析了。

所以 floorKey 方法就是传入一个 key，如果 TreeMap 存在该 key，则返回该 key，不存在则返回比该 key 稍小的上一个 key。如果传入的 key 是最小的 key，则返回 null。

而 TreeSet 的 floor 方法底层就是调用了 TreeMap 的 floorKey 方法。所以 TreeSet 的 floor 方法就是返回传入元素或者是传入元素的上一元素，不存在比传入元素更小的元素则返回 null。