java基础：TreeMap — 源码分析

其余更多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）

如上图，若是一个树有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. .....

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比较器进行排序。