Treemap&TreeSet源码分析

概述

注

TreeSet底层使用TreeMap实现，因次本文重点分析TreeMap

TreeMap 是一个有序的key-value集合，它是通过红黑树实现的。默认根据其键的自然顺序进行排序，或者根据创建映射时提供的 Comparator 进行排序，具体取决于使用的构造方法。

红黑树是一种近似平衡的二叉查找树，它能够确保任何一个节点的左右子树的高度差不会超过二者中较低那个的一倍。具体来说，红黑树是满足如下条件的二叉查找树(binary search tree):

1. 每个节点要么是红色，要么是黑色。
2. 根节点必须是黑色
3. 红色节点不能连续(也即是，红色节点的孩子和父亲都不能是红色)。
4. 对于每个节点，从该点至null(树尾端)的任何路径，都含有相同个数的黑色节点。

在树的结构发生改变时(插入或者删除操作)，往往会破坏上述条件3或条件4，需要通过调整使得查找树重新满足红黑树的约束条件。

方法解析

• 新增数据

  put(K key, V value)方法是将指定的key, value对添加到map里。该方法首先会对map做一次查找，看是否包含该元组，如果已经包含则直接返回，查找过程类似于getEntry()方法；如果没有找到则会在红黑树中插入新的entry，如果插入之后破坏了红黑树的约束条件，还需要进行调整(旋转，改变某些节点的颜色)。

  public V put(K key, V value) {
      Entry<K,V> t = root;
      if (t == null) {//如果root为null 说明是添加第一个元素 直接实例化一个Entry 赋值给root
          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;//如果root不为null，说明已存在元素 
      // split comparator and comparable paths
      Comparator<? super K> cpr = comparator; 
      if (cpr != null) { //如果比较器不为null 则使用比较器
          //找到元素的插入位置
          do {
              parent = t; //parent赋值
              cmp = cpr.compare(key, t.key);
              //当前key小于节点key 向左子树查找
              if (cmp < 0)
                  t = t.left;
              else if (cmp > 0)//当前key大于节点key 向右子树查找
                  t = t.right;
              else //相等的情况下 直接更新节点值
                  return t.setValue(value);
          } while (t != null);
      }
      else { //如果比较器为null 则使用默认比较器
          if (key == null)//如果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;
  }
  
  

  插入数据保证红黑树性质实现：

  private void fixAfterInsertion(Entry<K,V> x) {
      // 将新插入节点的颜色设置为红色
      x. color = RED;
      // while循环，保证新插入节点x不是根节点或者新插入节点x的父节点不是红色（这两种情况不需要调整）
      while (x != null && x != root && x. parent.color == RED) {
          // 如果新插入节点x的父节点是祖父节点的左孩子
          if (parentOf(x) == leftOf(parentOf (parentOf(x)))) {
              // 取得新插入节点x的叔叔节点
              Entry<K,V> y = rightOf(parentOf (parentOf(x)));
              // 如果新插入x的父节点是红色
              if (colorOf(y) == RED) {
                  // 将x的父节点设置为黑色
                  setColor(parentOf (x), BLACK);
                  // 将x的叔叔节点设置为黑色
                  setColor(y, BLACK);
                  // 将x的祖父节点设置为红色
                  setColor(parentOf (parentOf(x)), RED);
                  // 将x指向祖父节点，如果x的祖父节点的父节点是红色，按照上面的步奏继续循环
                  x = parentOf(parentOf (x));
              } else {
                  // 如果新插入x的叔叔节点是黑色或缺少，且x的父节点是祖父节点的右孩子
                  if (x == rightOf( parentOf(x))) {
                      // 左旋父节点
                      x = parentOf(x);
                      rotateLeft(x);
                  }
                  // 如果新插入x的叔叔节点是黑色或缺少，且x的父节点是祖父节点的左孩子
                  // 将x的父节点设置为黑色
                  setColor(parentOf (x), BLACK);
                  // 将x的祖父节点设置为红色
                  setColor(parentOf (parentOf(x)), RED);
                  // 右旋x的祖父节点
                  rotateRight( parentOf(parentOf (x)));
              }
          } else { // 如果新插入节点x的父节点是祖父节点的右孩子和上面的相似
              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;
  }
  

• 删除元素

  public V remove(Object key) {
          // 根据key查找到对应的节点对象
          Entry<K,V> p = getEntry(key);
          if (p == null)
              return null;
  
          // 记录key对应的value，供返回使用
          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) {
              // 查找p的替代节点
              Entry<K,V> s = successor (p);
              p. key = s.key ;
              p. value = s.value ;
              p = s;
          }
          Entry<K,V> replacement = (p. left != null ? p.left : p. right);
          if (replacement != null) { 
              // 将p的父节点拷贝给替代节点
              replacement. parent = p.parent ;
              // 如果替代节点p的父节点为空，也就是p为跟节点，则将replacement设置为根节点
              if (p.parent == null)
                  root = replacement;
              // 如果替代节点p是其父节点的左孩子，则将replacement设置为其父节点的左孩子
              else if (p == p.parent. left)
                  p. parent.left   = replacement;
              // 如果替代节点p是其父节点的左孩子，则将replacement设置为其父节点的右孩子
              else
                  p. parent.right = replacement;
              // 将替代节点p的left、right、parent的指针都指向空
              p. left = p.right = p.parent = null;
              // 如果替代节点p的颜色是黑色，则需要调整红黑树以保持其平衡
              if (p.color == BLACK)
                  fixAfterDeletion(replacement);
          } else if (p.parent == null) { // return if we are the only node.
              // 如果要替代节点p没有父节点，代表p为根节点，直接删除即可
              root = null;
          } else {
              // 如果p的颜色是黑色，则调整红黑树
              if (p.color == BLACK)
                  fixAfterDeletion(p);
              // 下面删除替代节点p
              if (p.parent != null) {
                  // 解除p的父节点对p的引用
                  if (p == p.parent .left)
                      p. parent.left = null;
                  else if (p == p.parent. right)
                      p. parent.right = null;
                  // 解除p对p父节点的引用
                  p. parent = null;
              }
          }
      }
  

• 查找

  public V get(Object key) {
          Entry<K,V> p = getEntry(key);
          return (p==null ? null : p. value);
  }
  
  final Entry<K,V> getEntry(Object key) {
      / 如果比较器为空，只是用key作为比较器查询
      if (comparator != null) 
          return getEntryUsingComparator(key);
      if (key == null)
          throw new NullPointerException();
      Comparable<? super K> k = (Comparable<? super K>) key;
      // 取得root节点
      Entry<K,V> p = root;
     //核心来了：从root节点开始查找，根据比较器判断是在左子树还是右子树
      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;
  }
  

总结

• 基于红黑树的数据结构实现。，关于红黑树的具体使用，后续再算法专题深入研究
• TreeMap不是线程安全的，如需线程安全可使用Collections#synchronizedSortedMap
• 不允许插入为Null的key，HashMap可以为空，注意区别
• 可实现Comparator自定义排序规则