TreeMap 還能排序？分析下源碼就明白了

時間 2019-12-11

標籤 treemap 還能排序分析源碼明白简体版

原文原文鏈接

Java 中的 Map 是一種鍵值對映射，又被稱爲符號表或字典的數據結構，一般使用哈希表來實現，但也可以使用二叉查找樹、紅黑樹實現。java

HashMap 基於哈希表，但迭代時不是插入順序
LinkedHashMap 擴展了 HashMap，維護了一個貫穿全部元素的雙向鏈表，保證按插入順序迭代
TreeMap 基於紅黑樹，保證鍵的有序性，迭代時按鍵大小的排序順序

這裏就來分析下 TreeMap 的實現。基於紅黑樹，就意味着結點的增刪改查都能在 O(lgn) 時間複雜度內完成，若是按樹的中序遍歷就能獲得一個按 鍵-key 大小排序的序列。node

在看本文以前，建議看一下《紅黑樹這個數據結構，讓你又愛又恨？看了這篇，妥妥的征服它》對紅黑樹的分析，理解了紅黑樹，你會發現 TreeMap 如此簡單。算法

基本結構

TreeMap 的繼承結構以下，其中包含了一些關鍵字段和方法：數據結構

其中，相關字段的意義是：框架

Comparator - 不爲空，那麼就用它維持 key-鍵的有序，不然使用 key-鍵的天然順序
size - 記錄樹中結點的個數
modCount - 記錄樹結構變化次數，用於迭代器的快速失敗

另外一個字段是 Entry<K,V> root ，它表示根結點，初始爲空，樹結點的結構定義以下：3d

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 葉子節點 的問題，用於避免在主算法中因檢查空葉子結點引發的混亂，方法以下：code

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): 刪除結點後的回調方法，從新平衡

這些方法基本上都能見名知意，其中有點繞的就是樹旋轉的代碼，代碼實現以下：blog

插入

結點的插入可能會打破紅黑樹的平衡，須要作旋轉和顏色變換的調整。假設待插入結點爲 N，P 是 N 的父結點，G 是 N 的祖父結點，U 是 N 的叔叔結點（即父結點的兄弟結點），那麼紅黑樹有如下幾種插入狀況：排序

N 是根結點，即紅黑樹的第一個結點
N 的父結點（P）爲黑色
P 是紅色的（不是根結點），它的兄弟結點 U 也是紅色的
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 的大小，遞歸的判斷插入左子樹仍是右子樹，比較複雜操做在於插入後從新平衡的調整，核心代碼以下：

刪除

結點的刪除也可能會打破紅黑樹的平衡，相比插入它的狀況更復雜，假設待刪除結點爲 M，若是有非葉子結點，稱爲 C，那麼有兩種比較簡單的刪除狀況：

M 爲紅色結點，那麼它必是葉子結點，直接刪除便可，由於若是它有一個黑色的非葉子結點，那麼就違反了性質5，經過 M 向左或向右的路徑黑色結點不等
M 是黑色而 C 是紅色，只須要讓 C 替換到 M 的位置，並變成黑色便可，或者說交換 C 和 M 的值，並刪除 C（就是第一個簡單的狀況）

這兩個狀況，本質都是刪除了一個紅色結點，不影響總體平衡，比較複雜的是 M 和 C 都是黑色的狀況，須要找一個結點填補這個黑色空缺。

結點 M刪除後它的位置上就變成了 NIL 隱形結點，爲了方便描述，這個結點記爲 N，P 表示 N 的父結點，S 表示 N 兄弟結點，S 若是存在左右孩子，分別使用 SL 和 SR 表示，那麼刪除就有如下幾種狀況：

N 是根結點 - 直接刪除便可
P 黑 S 紅 - 交換 P 和 S 的顏色，而後對 P 左旋轉
P 黑 S 黑 - 將 S 變成紅色，問題遞歸到父結點處理
P 紅 S 黑 - 將 S 變成紅色，刪除成功
P 顏色任意 S 黑 SL 紅 - 對 S 右旋轉，並交換 S 和 SL 的顏色，變成狀況6
P 顏色任意 S 黑，SR 紅 - 對 P 左旋轉，交換 P 和 S 的顏色，並將 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 的邏輯就和二叉查找樹同樣，主要就是把刪除任一結點的問題就簡化成：刪除一個最多隻有一個孩子的結點的狀況，而且全部的刪除操做都在葉子結點完成。若是刪除的是黑色結點，那麼就視狀況調整樹從新達到平衡，具體代碼以下: