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 也會用得遊刃有餘。

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章