TreeMap 还能排序?分析下源码就明白了

Java 中的 Map 是一种键值对映射,又被称为符号表字典的数据结构,通常使用哈希表来实现,但也可使用二叉查找树红黑树实现。

  • HashMap 基于哈希表,但迭代时不是插入顺序
  • LinkedHashMap 扩展了 HashMap,维护了一个贯穿所有元素的双向链表,保证按插入顺序迭代
  • TreeMap 基于红黑树,保证有序性,迭代时按大小的排序顺序

这里就来分析下 TreeMap 的实现。基于红黑树,就意味着结点的增删改查都能在 O(lgn) 时间复杂度内完成,如果按树的中序遍历就能得到一个按 键-key 大小排序的序列。

在看本文之前,建议看一下《红黑树这个数据结构,让你又爱又恨?看了这篇,妥妥的征服它》对红黑树的分析,理解了红黑树,你会发现 TreeMap 如此简单。

基本结构

TreeMap 的继承结构如下,其中包含了一些关键字段和方法:

class.png

其中,相关字段的意义是:

  • Comparator - 不为空,那么就用它维持 key-键 的有序,否则使用 key-键 的自然顺序
  • size - 记录树中结点的个数
  • modCount - 记录树结构变化次数,用于迭代器的快速失败

另一个字段是 Entry<K,V> root ,它表示根结点,初始为空,树结点的结构定义如下:

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

TreeMap 是按照算法导论(CLR)的描述实现的,但略有不同,它没有使用隐形叶子结点 NIL,而是定义了一组访问方法来正确处理 NULL 叶子节点 的问题,用于避免在主算法中因检查空叶子结点引起的混乱,方法如下:

  • colorOf(Entry<K,V> p): 返回结点颜色,如果为空返回黑色
  • parentOf(Entry<K,V> p): 返回父结点的引用,根结点则返回 null
  • setColor(Entry<K,V> p, boolean c): 设置结点颜色
  • leftOf(Entry<K,V> p): 返回左孩子结点
  • rightOf(Entry<K,V> p): 返回右孩子结点
  • rotateLeft(Entry<K,V> p): 将结点 P 左旋转
  • rotateRight(Entry<K,V> p): 将结点 P 右旋转
  • fixAfterInsertion(Entry<K,V> x): 插入结点后的回调方法,重新平衡
  • fixAfterDeletion(Entry<K,V> x): 删除结点后的回调方法,重新平衡

这些方法基本上都能见名知意,其中有点绕的就是树旋转的代码,代码实现如下:

rotate.png

插入

结点的插入可能会打破红黑树的平衡,需要做旋转和颜色变换的调整。假设待插入结点为 NPN 的父结点,GN 的祖父结点,UN 的叔叔结点(即父结点的兄弟结点),那么红黑树有以下几种插入情况:

  1. N 是根结点,即红黑树的第一个结点
  2. N 的父结点(P)为黑色
  3. P红色的(不是根结点),它的兄弟结点 U 也是红色
  4. P红色,而 U黑色
    4.1 P 左(右)孩子 N 右(左)孩子
    4.2 P 左(右)孩子 N 左(右)孩子

以上情况的分析可查看本文开头的文章链接,现在来看下 TreeMap 的 put 方法的实现:

public V put(K key, V value) {
  Entry<K,V> t = root;
  // 情况 1 - 空树,直接插入作为根结点
  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) { // 使用 comparator 比较大小
    do { // 根据 key 的大小找到插入位置
      parent = t;
      cmp = cpr.compare(key, t.key);
      if (cmp < 0) t = t.left;
      else if (cmp > 0) t = t.right;
      else // 如果有相等的 key 直接设置 value 并返回 
        return t.setValue(value);
    } while (t != null);
  }
  else {// 使用 key 的自然顺序
    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;
}

put 方法比较简单,就是根据 key 的大小,递归的判断插入左子树还是右子树,比较复杂操作在于插入后重新平衡的调整,核心代码如下:

fix-insert.png

删除

结点的删除也可能会打破红黑树的平衡,相比插入它的情况更复杂,假设待删除结点为 M,如果有非叶子结点,称为 C,那么有两种比较简单的删除情况:

  1. M 为红色结点,那么它必是叶子结点,直接删除即可,因为如果它有一个黑色的非叶子结点,那么就违反了性质5,通过 M 向左或向右的路径黑色结点不等
  2. M 是黑色而 C 是红色,只需要让 C 替换到 M 的位置,并变成黑色即可,或者说交换 CM 的值,并删除 C(就是第一个简单的情况)

这两个情况,本质都是删除了一个红色结点,不影响整体平衡,比较复杂的是 MC 都是黑色的情况,需要找一个结点填补这个黑色空缺

结点 M删除后它的位置上就变成了 NIL 隐形结点,为了方便描述,这个结点记为 NP 表示 N 的父结点,S 表示 N 兄弟结点,S 如果存在左右孩子,分别使用 SLSR 表示,那么删除就有以下几种情况:

  1. N 是根结点 - 直接删除即可
  2. PS 红 - 交换 PS 的颜色,然后对 P 左旋转
  3. PS 黑 - 将 S 变成红色,问题递归到父结点处理
  4. PS 黑 - 将 S 变成红色,删除成功
  5. P 颜色任意 SSL 红 - 对 S 右旋转,并交换 SSL 的颜色,变成情况6
  6. P 颜色任意 S 黑,SR 红 - 对 P 左旋转,交换 PS 的颜色,并将 SR 变成黑色

针对这些情况,TreeMap 进行了实现:

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

  V oldValue = p.value;
  deleteEntry(p); // 删除结点
  return oldValue;
}
private void deleteEntry(Entry<K,V> p) {
  modCount++;
  size--;
  // 如果 p 有两个孩子结点,转成删除最多有一个孩子的结点的情况
  // 这里查找的是 p 的后继结点,也就是右子树值最小的结点
  if (p.left != null && p.right != null) {
    Entry<K,V> s = successor(p); // 查找后继结点
    // 复制后继结点的 key 和 value 到 p
    p.key = s.key;
    p.value = s.value;
    p = s; // 将 p 指向这个右子树值最小的结点
  } // p has 2 children

  // 此时删除的 p 要么是叶子结点,要么只有一个左或右孩子
  Entry<K,V> replacement = (p.left != null ? p.left : p.right);

  if (replacement != null) { // 有孩子结点
    // 有一个左或右孩子,使用这个孩子结点替换它的父结点 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;

    // Null out links so they are OK to use by fixAfterDeletion.
    // 删除结点 p,也就是断开所有的链接
    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;// 情况1,删除后变成空树
  } else {//No children. Use self as phantom replacement and unlink.
    // 删除的是叶子结点,那么删除 p 就是用它的隐形 NIL 叶子结点替换
    // 它,这里将它自己看做隐形的叶子结点
    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;
    }
  }
}

deleteEntry 的逻辑就和二叉查找树一样,主要就是把删除任一结点的问题就简化成:删除一个最多只有一个孩子的结点的情况,并且所有的删除操作都在叶子结点完成。如果删除的是黑色结点,那么就视情况调整树重新达到平衡,具体代码如下:

fix-delete.png

查找

就像二分查找那样,TreeMap 也能在 ~lgN 次比较内结束查找,并且针对 键-key 提供了丰富的查询 API,

  • get(Object key) - 返回等于给定键的结点
  • floorEntry(K key) - 返回小于或等于给定键的结点中键最大的结点
  • ceilingEntry(K key) - 返回大于或等于给定键的结点中键最小的结点
  • higherEntry(K key) - 返回严格大于给定键的结点中键最小的结点
  • lowerEntry(K key) - 返回严格小于给定键的结点中键最大的结点

上面这些方法比较简单,可自行查看源码。另外,还有两个比较特殊的方法,它们用来查询指定结点在树中序遍历序列中的前驱和后继结点,在中序遍历序列中:

  • 前驱结点也就是左子树值最大的结点
  • 后继结点也就是右子树值最小的结点

serach-node.png

遍历

遍历也是一个高频操作,在 Java 集合框架体系中,基本都是采用迭代器 Iterator 来实现,TreeMap 也是如此,它提供了对和对的迭代器。

TreeMap 迭代器最终的逻辑实现是在 PrivateEntryIterator 类中,默认按键的正序输出,它也提供了一个逆序输出的迭代器 DescendingKeyIterator。

具体代码不在贴出,比较简单,值得注意的就是上一节介绍的查找前驱和后继结点的两个方法,遍历常用 API 有:

  • entrySet() - 返回一个遍历所有结点的 Set 集合
  • keySet() - 返回一个遍历所有的 Set 集合
  • values() - 返回一个遍历所有的 Set 集合

小结

分析 TreeMap 的源码之前,一定要去分析红黑树的原理,然后在看它的源码,相信理论与实践相结合,掌握红黑树不在话下,TreeMap 也会用得游刃有余。

posted @ 2019-07-26 09:42 小创 阅读(...) 评论(...) 编辑 收藏