java基礎:TreeMap — 原始碼分析

Hiway發表於2019-01-02

其他更多java基礎文章:
java基礎學習(目錄)


概述

TreeMap的底層資料結構就是一個紅黑樹。關於紅黑樹的知識可以檢視演算法--我的紅黑樹學習過程
TreeMap的特點就是儲存的時候是根據鍵Key來進行排序的。其順序與新增順序無關,該順序根據key的自然排序進行排序或者根據構造方法中傳入的Comparator比較器進行排序。自然排序要求key需要實現Comparable介面。

資料結構和基礎欄位

    //比較器,若無則按Key的自然排序
    private final Comparator<? super K> comparator;
    //樹根結點
    private transient Entry<K,V> root;
    //樹節點個數
    private transient int size = 0;
    //用於判斷資料是否變化
    private transient int modCount = 0;
複製程式碼

Entry<K,V>表示紅黑樹的一個結點,既然是紅黑樹,那麼每個節點中除了Key-->Value對映之外,必然儲存了紅黑樹節點特有的一些內容

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;//黑色表示為true,紅色為false
}
複製程式碼

方法細節

構造方法

public TreeMap() {
    comparator = null;
}

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


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

public void putAll(Map<? extends K, ? extends V> map) {
  int mapSize = map.size();
  //判斷map是否SortedMap,不是則採用AbstractMap的putAll
  if (size==0 && mapSize!=0 && map instanceof SortedMap) {
    Comparator<?> c = ((SortedMap<?,?>)map).comparator();
    //同為null或者不為null,型別相同,則進入有序map的構造
    if (c == comparator || (c != null && c.equals(comparator))) {
      ++modCount;
      try {
        buildFromSorted(mapSize, map.entrySet().iterator(),
                        null, null);
      } catch (java.io.IOException cannotHappen) {
      } catch (ClassNotFoundException cannotHappen) {
      }
      return;
    }
  }
  super.putAll(map);
}
複製程式碼

都比較簡單,我們主要關注一下buildFromSorted方法和computeRedLevel方法。TreeMap主要通過這兩個方法在初始化的時候構造一個簡單的紅黑樹。

private static int computeRedLevel(int sz) {
  int level = 0;
  for (int m = sz - 1; m >= 0; m = m / 2 - 1)
    level++;
  return level;
}
複製程式碼

computeRedLevel方法是計算當前結點數的完全二叉樹的層數。或者說,著色紅色結點的層數。
TreeMap是如何構造紅黑樹的呢,簡單來說,就是把當前的結點按照完全二叉樹的結構來排列,此時,最下層的符合二叉樹又未滿足滿二叉樹 的那一排結點,就全部設為紅色,這樣就滿足紅黑樹的條件了。(TreeMap中第一層根結點層數為0)

java基礎:TreeMap — 原始碼分析
如上圖,如果一個樹有9個節點,那麼我們構造紅黑樹的時候,只要把前面3層的結點都設定為黑色,第四層的節點設定為紅色,則構造完的樹,就是紅黑樹。而實現的關鍵就是找到要構造樹的完全二叉樹的層數

瞭解了上面的原理,後面就簡單了,接著來看buildFromSorted方法:

/**
* level: 當前樹的層數,注意:是從0層開始
* lo: 子樹第一個元素的索引
* hi: 子樹最後一個元素的索引
* redLevel: 上述紅節點所在層數
* it: 傳入的map的entries迭代器
* str: 如果不為空,則從流裡讀取key-value
* defaultVal:不為空,則value都用這個值
*/
@SuppressWarnings("unchecked")
private final Entry<K,V> buildFromSorted(int level, int lo, int hi,
                                         int redLevel,
                                         Iterator<?> it,
                                         java.io.ObjectInputStream str,
                                         V defaultVal)
  throws  java.io.IOException, ClassNotFoundException {
  // hi >= lo 說明子樹已經構造完成
  if (hi < lo) return null;
  // 取中間位置,無符號右移,相當於除2
  int mid = (lo + hi) >>> 1;
  Entry<K,V> left  = null;
  //遞迴構造左結點
  if (lo < mid)
    left = buildFromSorted(level+1, lo, mid - 1, redLevel,
                           it, str, defaultVal);
  K key;
  V value;
  //遞迴完左子樹後,迭代器的下一個結點就是每棵樹或子樹的根結點,所以此時獲取的key,value就是樹根結點的key,value
  if (it != null) {
    if (defaultVal==null) {
      Map.Entry<?,?> entry = (Map.Entry<?,?>)it.next();
      key = (K)entry.getKey();
      value = (V)entry.getValue();
    } else {
      key = (K)it.next();
      value = defaultVal;
    }
  // 通過流來讀取key, value
  } else {
    key = (K) str.readObject();
    value = (defaultVal != null ? defaultVal : (V) str.readObject());
  }
  //構建結點
  Entry<K,V> middle =  new Entry<>(key, value, null);
  // 這裡是判斷該節點是否是最下層的葉子結點。
  if (level == redLevel)
    middle.color = RED;
  //如果存在的話,設定左結點,
  if (left != null) {
    middle.left = left;
    left.parent = middle;
  }
  // 遞迴構造右結點
  if (mid < hi) {
    Entry<K,V> right = buildFromSorted(level+1, mid+1, hi, redLevel,
                                       it, str, defaultVal);
    middle.right = right;
    right.parent = middle;
  }
  return middle;
}
複製程式碼

舉個簡單例子,我們存入key為1,2,3,4,5,6的treeMap,程式碼如下:

    public static void main(String[] args) {
        TreeMap treeMap = new TreeMap();
        treeMap.put(1,1);
        treeMap.put(2,2);
        treeMap.put(3,3);
        treeMap.put(4,4);
        treeMap.put(5,5);
        treeMap.put(6,6);
        TreeMap map2 = new TreeMap(treeMap);
    }
複製程式碼

我們通過debug可以得出它最後構造的紅黑樹如下,結點外和連線旁的數字表示構造的順序,如:

  1. 構造結點1
  2. 構造出紅結點2
  3. 將結點1的right設為結點2,結點2的parent設為結點1
  4. .....

java基礎:TreeMap — 原始碼分析

get方法

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

final Entry<K,V> getEntry(Object key) {
    // Offload comparator-based version for sake of performance
    //  comparator 這個是 個成員變數 外部設定特定的 比較器 有就用這個  這個變數 可以初始化的時候 放進去
    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;
    // 利用比較器的特性開始比較大小  相同 return 小於 從左子樹開始 大了 從右子樹開始 
    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;
}
複製程式碼

get方法比較簡單,我們就不多講解了。

put方法

    //新增元素
    public V put(K key, V value) {
        //記錄根節點
        Entry<K,V> t = root;
        //如果根節點為空,該元素設定為root
        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) {
            //迴圈比較並確定元素插入的位置(找父親節點)
            do {
                //記錄根節點
                parent = t;
                //將當前節點和根節點元素比較
                cmp = cpr.compare(key, t.key);
                //待插入key小於當前元素key,查詢左邊
                if (cmp < 0)
                    t = t.left;
                //待插入key大於當前元素key,查詢右邊
                else if (cmp > 0)
                    t = t.right;
                //相等,替換
                else
                    return t.setValue(value);
            } while (t != null);
        }
        //比較器為null
        else {
            //TreeMap元素,key不能為null
            if (key == null)
                throw new NullPointerException();
            @SuppressWarnings("unchecked")
            //key需要實現Comparable介面
            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);
        //如果待插入元素的key值小於父節點的key值,父節點左邊插入
        if (cmp < 0) {
            parent.left = e;
        }
        //如果待插入元素的key值大於父節點的key值,父節點右邊插入
        else {
            parent.right = e;
        }
        //對紅黑樹進行重新平衡
        fixAfterInsertion(e);
        size++;
        modCount++;
        return null;
    }
複製程式碼

put方法也不難,根據排序的大小,去左子樹和右子樹裡查。主要是看fixAfterInsertion這個方法。在看這個方法之前,必須要了解紅黑樹的插入的幾種平衡情況,可以看紅黑樹詳細分析,看了都說好

/** From CLR */
    private void fixAfterInsertion(Entry<K,V> x) {
        // 符合一般規則 先插入的節點變為紅色
        x.color = RED;
            // 若“父節點存在,並且父節點的顏色是紅色”  文章中的情況三  將一直往上朔
        while (x != null && x != root && x.parent.color == RED) {
        
        
            //X的父節點是x祖父節點的左子樹
            if (parentOf(x) == leftOf(parentOf(parentOf(x)))) {
            
            
                // X 的 叔叔節點
                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 = parentOf(parentOf(x)); // 改變當前節點位置  將“祖父節點”設為“當前節點
                    
                    
                } else { // 情況四  或 五  叔叔是黑色,
                
                    if (x == rightOf(parentOf(x))) {  //情況四  
                        x = parentOf(x);
                        rotateLeft(x);          //左旋轉
                    }
                    setColor(parentOf(x), BLACK);  //情況五
                    setColor(parentOf(parentOf(x)), RED);
                    rotateRight(parentOf(parentOf(x)));
                }
                
                
                
            } else {
                //叔叔節點
                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 = parentOf(parentOf(x));
                } else {// 情況四  或 五  叔叔是黑色,
                    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; 
    }
複製程式碼

可以看出,插入平衡主要就是分為父節點為祖父節點左右子樹的情況分別判斷旋轉,以及紅黑樹詳細分析,看了都說好文中情況三,情況四,情況五的判斷。

remove方法

在講remove方法之前,先要了解一個尋找當前節點後繼的方法successor(Entry<K,V> t),因為在刪除方法中,如果刪除的不是最底層節點,需要尋找它的後繼節點來替換刪除。請先了解紅黑樹刪除的知識點,可以檢視演算法--我的紅黑樹學習過程

static <K,V> TreeMap.Entry<K,V> successor(Entry<K,V> t) {
    
        if (t == null)
            return null;
        // 如果right 不為空 往左
        else if (t.right != null) {
            //在t的右子樹的節點p的左子樹中迴圈查詢
            Entry<K,V> p = t.right;
            // while 迴圈找到中序後繼結點  一直往左找
            while (p.left != null)
                p = p.left;
            return p;
        } else {
            
            Entry<K,V> p = t.parent;
            
            Entry<K,V> ch = t;
            
           // while 迴圈找到中序後繼結點  一直往右找
            while (p != null && ch == p.right) {
                ch = p;
                p = p.parent;
            }
            return p;
        }
    }
複製程式碼

瞭解了successor方法後,我們來看下remove方法

    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的左右子樹都非空,則尋找後繼節點
        if (p.left != null && p.right != null) {
            Entry<K,V> s = successor(p);  //找出 中序後繼 節點
            p.key = s.key;
            p.value = s.value;
            p = s;
        } // p has 2 children

        
        Entry<K,V> replacement = (p.left != null ? p.left : p.right);
        
        // replacement是刪除節點的左子樹或右子樹,如果不為null,則表示刪除點p有一個子樹。
        //此時刪除節點p,只需要把節點p的子樹提上來到節點p的位置就可以了。
        if (replacement != null) {

            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.left = p.right = p.parent = null;

            // Fix replacement
            if (p.color == BLACK)
                fixAfterDeletion(replacement);
                
                
        } else if (p.parent == null) { // 只有一個節點
            root = null;
            
        } else { //左右子樹 都為空
            if (p.color == BLACK)
                fixAfterDeletion(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;
            }
        }
    }
複製程式碼

可以看到,刪除節點的邏輯並不複雜,重要的還是刪除節點後的平衡節點方法fixAfterDeletion

private void fixAfterDeletion(Entry<K,V> x) {
  while (x != root && colorOf(x) == BLACK) {
    //x是左結點且為黑色
    if (x == leftOf(parentOf(x))) {
      //獲取兄弟右節點
      Entry<K,V> sib = rightOf(parentOf(x));
      //① D黑,S紅
      if (colorOf(sib) == RED) {
        setColor(sib, BLACK);
        setColor(parentOf(x), RED);
        rotateLeft(parentOf(x));
        sib = rightOf(parentOf(x));
      }
      //② D黑,S黑,SL黑,SR黑
      if (colorOf(leftOf(sib))  == BLACK &&
          colorOf(rightOf(sib)) == BLACK) {
        setColor(sib, RED);
        //此時將x設為父節點,因為不管P是黑是紅,最後都會設為黑。
        x = parentOf(x);
      //sib子節點不全為黑
      } else {
        //③ sib右子節點為黑色,D黑,S黑,SL紅
        if (colorOf(rightOf(sib)) == BLACK) {
          setColor(leftOf(sib), BLACK);
          setColor(sib, RED);
          rotateRight(sib);
          sib = rightOf(parentOf(x));
        }
        // ④D黑,S黑,SR紅
        setColor(sib, colorOf(parentOf(x)));
        setColor(parentOf(x), BLACK);
        setColor(rightOf(sib), BLACK);
        rotateLeft(parentOf(x));
        x = root;
      }
    // 對稱
    } else {
      Entry<K,V> sib = leftOf(parentOf(x));

      if (colorOf(sib) == RED) {
        setColor(sib, BLACK);
        setColor(parentOf(x), RED);
        rotateRight(parentOf(x));
        sib = leftOf(parentOf(x));
      }

      if (colorOf(rightOf(sib)) == BLACK &&
          colorOf(leftOf(sib)) == BLACK) {
        setColor(sib, RED);
        x = parentOf(x);
      } else {
        if (colorOf(leftOf(sib)) == BLACK) {
          setColor(rightOf(sib), BLACK);
          setColor(sib, RED);
          rotateLeft(sib);
          sib = leftOf(parentOf(x));
        }
        setColor(sib, colorOf(parentOf(x)));
        setColor(parentOf(x), BLACK);
        setColor(leftOf(sib), BLACK);
        rotateRight(parentOf(x));
        x = root;
      }
    }
  }
  setColor(x, BLACK);
}
複製程式碼

程式碼中的註釋配合著演算法--我的紅黑樹學習過程對照檢視,就非常簡單了,我已經將紅黑樹刪除的幾種情況分別註釋了。

總結

  1. TreeMap的底層資料結構就是一個紅黑樹,增刪改查和統計相關的操作的時間複雜度都為 O(logn)。所以瞭解紅黑樹的資料結構和邏輯就非常重要。
  2. TreeMap的特點就是儲存的時候是根據鍵Key來進行排序的。
  3. key的排序分為自然排序和比較器排序。自然排序要求key需要實現Comparable介面,比較器排序需要在構造方法中傳入的Comparator比較器進行排序。

相關文章