TreeMap实现原理及源码浅析

一、简介

TreeMap是一个线程不安全，有序的键值对集合，因为TreeMap实现了SotredMap接口。底层是一个红黑树的数据结构，每个key-value都作为一个红黑树的节点。如果在调用TreeMap的构造函数时没有指定比较器，则根据key执行自然排序。这点会在接下来的代码中做说明，如果指定了比较器则按照比较器来进行排序。
此外，TreeMap的Key是不能为null的，因为要实现排序。而Value是允许为null的。

二、数据结构

/**
* TreeMap
*/
public class TreeMap<K,V>
    extends AbstractMap<K,V>
    implements NavigableMap<K,V>, Cloneable, java.io.Serializable
{
    // 创建Map传进去的比较器，不传为null
    private final Comparator<? super K> comparator;
	// 存储根节点
    private transient TreeMapEntry<K,V> root;
    // 大小
    private transient int size = 0;
    // 修改次数
    private transient int modCount = 0;
}

接下来看看TreeMapEntry的数据结构，形成红黑树和存储最重要的实体类

/**
* TreeMap的静态内部类TreeMapEntry
*/
static final class TreeMapEntry<K,V> implements Map.Entry<K,V> {
		// 保存我们存储的Key
        K key; 
        // 保存我们存储的Value
        V value; 
        // 当前节点的左子节点
        TreeMapEntry<K,V> left; 
        // 当前节点右子节点
        TreeMapEntry<K,V> right; 
        // 当前节点父节点
        TreeMapEntry<K,V> parent; 
        // 红黑树的节点颜色
        boolean color = BLACK; 
 		// 构造方法，KV及父节点
        TreeMapEntry(K key, V value, TreeMapEntry<K,V> parent) {
            this.key = key;
            this.value = value;
            this.parent = parent;
        }
}

TreeMap构造方法如下：

/**
* TreeMap构造方法
*/
// 无参构造方法，
public TreeMap() {
    comparator = null;
}

// 带有比较器comparator的构造方法
public TreeMap(Comparator<? super K> comparator) {
    this.comparator = comparator;
}

// 原集合需要添加到新集合的构造方法
public TreeMap(Map<? extends K, ? extends V> m) {
     comparator = null;
     putAll(m);
}

// 参数为SortedMap的构造方法
// 采用m的比较器排序
public TreeMap(SortedMap<K, ? extends V> m) {
     comparator = m.comparator();
     try {
         buildFromSorted(m.size(), m.entrySet().iterator(), null, null);
     } catch (java.io.IOException cannotHappen) {
     } catch (ClassNotFoundException cannotHappen) {
     }
}

TreeMap提供了四个构造方法，实现了方法的重载。无参构造方法中比较器的值为null，采用自然排序的方法，如果指定了比较器则称之为定制排序。

• 自然排序：TreeMap的所有key必须实现Comparable接口，所有的key都是同一个类的对象。
• 定制排序：创建TreeMap对象传入了一个Comparator对象，该对象负责对TreeMap中所有的key进行排序，采用定制排序不要求Map的key实现Comparable接口。

三、put方法

/**
* TreeMap类的put方法
*/
public V put(K key, V value) {
	 // 获取成员变量保存的根节点，赋值给t
     TreeMapEntry<K,V> t = root;
     // 根节点还没有
     if (t == null) {
         compare(key, key); // type (and possibly null) check
		 // 创建一个根节点，赋值给成员变量root，根节点的parent必然是null啦
         root = new TreeMapEntry<>(key, value, null);
         size = 1;
         modCount++;
         return null;
     }
     int cmp;
     TreeMapEntry<K,V> parent;
     // 获取创建TreeMap时设置的比较器
     Comparator<? super K> cpr = comparator;
     if (cpr != null) {
     	 // 有定制的比较器
     	 // 通过do-while循环的方式，借助定制的比较器，来找出需要插入的这个key的parent是谁
         do {
             parent = t;
             cmp = cpr.compare(key, t.key);
             if (cmp < 0)
                 t = t.left;
             else if (cmp > 0)
                 t = t.right;
             else
                 return t.setValue(value);
         } while (t != null);
     }
     else {
     	 // 没有定制的比较器，自然排序的情况
     	 // 这种情况是不允许我们插入null的key
         if (key == null)
             throw new NullPointerException();
         @SuppressWarnings("unchecked")
         // 获取key本身的比较器Comparable
         // 这也是为什么上面说的，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);
     }
     // 找到要插入的KV的parent后，创建一个TreeMapEntry对象存储这些数据信息
     TreeMapEntry<K,V> e = new TreeMapEntry<>(key, value, parent);
     // 根据和父节点的比较结果，是放在父节点的左边还是右边
     if (cmp < 0)
         parent.left = e;
     else
         parent.right = e;
     // 这个方法就比较复杂了，新插入节点后重新调整红黑树
     fixAfterInsertion(e);
     size++;
     modCount++;
     return null;
}

在树的基本结构那一篇博客里，我们分析过红黑树的插入，提前了解有助于更好理解下面的代码。
红黑树插入过程

/**
* TreeMap插入数据后，红黑树重新调整的方法
*/
private void fixAfterInsertion(TreeMapEntry<K,V> x) {
	 // 插入的节点默认是红色的
     x.color = RED;
	 //（1）新添加节点N为根：涂黑完事（因为根节点必须为黑）
	 //（2）新添加的节点的父节点是黑的：啥事不用管，添加进来就行
     while (x != null && x != root && x.parent.color == RED) {
     	 // 这里面就是处理我那篇博客里的（3）和（4.a）（4.b）的情况
     	 // 判断x的节点的父节点位置，是否属于左孩子
         if (parentOf(x) == leftOf(parentOf(parentOf(x)))) {
         	 // 获取叔叔节点
             TreeMapEntry<K,V> y = rightOf(parentOf(parentOf(x)));
             if (colorOf(y) == RED) {
             	 // 属于情况（3）
             	 // （3）新添加节点的父红和叔红：父/叔涂黑，祖父涂红，然后把祖父当成新的平衡节点递归处理
                 setColor(parentOf(x), BLACK);
                 setColor(y, BLACK);
                 setColor(parentOf(parentOf(x)), RED);
                 x = parentOf(parentOf(x));
             } else {
              	 // 父红，叔黑	
                 if (x == rightOf(parentOf(x))) {
                 	 // （4.a）的情况，不在同一边，先旋转扭转
                     x = parentOf(x);
                     rotateLeft(x);
                 }
                 // （4.b）的情况，在同一边，然后以祖父为中心旋转
                 setColor(parentOf(x), BLACK);
                 setColor(parentOf(parentOf(x)), RED);
                 // 右旋
                 rotateRight(parentOf(parentOf(x)));
             }
         } else {
         	 // 父节点是属于右孩子的情况
         	 // 获取叔叔
             TreeMapEntry<K,V> y = leftOf(parentOf(parentOf(x)));
             if (colorOf(y) == RED) {
             	 // 叔叔是红的
             	 // 情况（3）
             	 // （3）新添加节点的父红和叔红：父/叔涂黑，祖父涂红，然后把祖父当成新的平衡节点递归处理
                 setColor(parentOf(x), BLACK);
                 setColor(y, BLACK);
                 setColor(parentOf(parentOf(x)), RED);
                 x = parentOf(parentOf(x));
             } else {
             	 // 父红，叔黑
                 if (x == leftOf(parentOf(x))) {
                 	 // （4.a）的情况，不在同一边，先旋转扭转
                     x = parentOf(x);
                     rotateRight(x);
                 }
                 // （4.b）的情况，在同一边，然后以祖父为中心旋转
                 setColor(parentOf(x), BLACK);
                 setColor(parentOf(parentOf(x)), RED);
                 // 左旋
                 rotateLeft(parentOf(parentOf(x)));
             }
         }
     }
     root.color = BLACK;
}