TreeMap jdk1.8源码分析

TreeMap特性：

• 基于红黑树实现的有序Map集合

• 不允许key值为空 (有一种情况允许key为null，设置了comparator，且comparator的compare方法，没有取空key进行处理，但是这种情况几乎没有意义)

public class TreeMap<K,V>
    extends AbstractMap<K,V>
    implements NavigableMap<K,V>, Cloneable, java.io.Serializable {
 	 //外部比较器，可以为null
    private final Comparator<? super K> comparator;
    
  	//红黑树的根节点
    private transient Entry<K,V> root;
    
	//红黑树的节点个数
    private transient int size = 0;

   //默认外部比较器为null
    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);
    }

    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) {
        }
    }
}

TreeMapEntry节点

红黑树的节点TreeMapEntry（key、 value、左节点left、右节点right、父节点parent、颜色color）

 static final class TreeMapEntry<K,V> implements Map.Entry<K,V> {
        K key;
        V value;
        TreeMapEntry<K,V> left;
        TreeMapEntry<K,V> right;
        TreeMapEntry<K,V> parent;
        boolean color = BLACK;

        Entry(K key, V value, TreeMapEntry<K,V> parent) {
            this.key = key;
            this.value = value;
            this.parent = parent;
        }

	 	public V setValue(V value) {
            V oldValue = this.value;
            this.value = value;
            return oldValue;
        }
    }

插入、修改元素put(key, value)

• 如果root为null，则新建Entry节点，作为root节点，插入成功
• 如果comparator不为空，从root节点开始遍历比较，comparator.compare(o1, o2)
• 如果comparator为空，从root节点开始遍历比较，o1.compareTo(o2) (如果key没有实现comparable接口，会报classCastException)
• • 如果key小于root节点的key，取root节点的左子节点作为新的root节点，继续和key比较
• • 如果key大于root节点的key，取root节点的右子节点作为新的root节点，继续和key比较
• • 如果key等于root节点的key，则用value值替换root节点的旧value值，然后返回root节点的旧value值
• • 直到遍历结束，也就是遍历到了尾部
• • • 新建一个叶子节点e = new Entry(key,value, parent)
• • • 如果key小于parent的key，则添加e为parent的左叶子节点，parent.left = e
• • • 如果key 大于parent的key，则添加e为parent的右叶子节点，parent.right = e

public V put(K key, V value) {
        Entry<K,V> t = root; 
        if (t == null) { //若这是一个空的红黑树结构
            compare(key, key); // key值校验，必须实现comparable接口或者指定comparator
            root = new Entry<>(key, value, null); //新建一个root节点
            size = 1;
            modCount++;
            return null; //返回，插入成功
        }
        int cmp;
        Entry<K,V> parent;
        // 区分比较 comparator 和 comparable 
        Comparator<? super K> cpr = comparator;
        if (cpr != null) {
            do { //comparator不为空
                parent = t;
                cmp = cpr.compare(key, t.key);
                if (cmp < 0) //如果key小于root节点的key
                    t = t.left; //取root节点的左子节点作为新的root节点，继续和key比较
                else if (cmp > 0) //如果key大于root节点的key
                    t = t.right; //取root节点的右子节点作为新的root节点，继续和key比较
                else //如果key等于root节点的key
                    return t.setValue(value); //则用newValue替换oldValue值，然后返回oldValue
            } while (t != null); //直到遍历到尾部
        } else {
            if (key == null) //不允许空键
                throw new NullPointerException();
            //没有指定comparator，则key必须实现comparable接口，否则会报classCastException
                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);
        }
        //遍历到尾部，新建一个叶子节点e
        Entry<K,V> e = new Entry<>(key, value, parent);
        if (cmp < 0) //如果key小于parent的key，则添加e为parent的左叶子节点
            parent.left = e;
        else //如果key大于parent的key，则添加e为parent的右叶子节点
            parent.right = e;
        fixAfterInsertion(e); //对这棵树进行调整、平衡
        size++;
        modCount++;
        return null;
    }

上面的do{}代码仅仅是排序二叉树的核心算法,真正的实现红黑树依靠的是fixAfterInsertion(e)方法

红黑树是一棵平衡排序二叉树，普通的排序二叉树可能会出现失衡的情况，而fixAfterInsertion就是调整结构实现红黑树 : 左旋、右旋、着色三个基本操作 。

查找元素get(key)

• 如果comparator不为空，从root节点开始遍历比较，comparator.compare(o1, o2)

• 如果comparator为空，从root节点开始遍历比较，o1.compareTo(o2) (如果key没有实现comparable接口，会报classCastException)

  • • 如果key小于root节点的key，取root节点的左子节点作为新的root节点，继续和key比较
  • • 如果key大于root节点的key，取root节点的右子节点作为新的root节点，继续和key比较
  • • 如果key等于root节点的key，则匹配成功，返回这个root节点
  • • 直到遍历结束，返回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) {
        // 如果comparator不为空，按comparator的方式，从root节点开始遍历比较
        if (comparator != null)
            return getEntryUsingComparator(key);
      //如果comparator为空，按comparable的方式，从root节点开始遍历比较
        Comparable<? super K> k = (Comparable<? super K>) key;
        Entry<K,V> p = 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; //遍历结束依然没匹配上，返回null
    }

//按comparator的方式，从root节点开始遍历比较
 final Entry<K,V> getEntryUsingComparator(Object key) {
        K k = (K) key;
        Comparator<? super K> cpr = comparator;
        if (cpr != null) {
            Entry<K,V> p = root;
            while (p != null) {
                int cmp = cpr.compare(k, p.key);
                if (cmp < 0)
                    p = p.left;
                else if (cmp > 0)
                    p = p.right;
                else
                    return p;//匹配上，返回节点
            }
        }
        return null;//遍历结束依然没匹配上，返回null
    }

删除元素

• 根据key定位到节点p，p 为null，则不存在此元素，直接返回null
• 如果节点p既有左节点，又有右节点

 public V remove(Object key) {
        Entry<K,V> p = getEntry(key); //根据key定位到节点p
        if (p == null) //p 为null，则不存在此元素
            return null;

        V oldValue = p.value;
        deleteEntry(p);
        return oldValue;
    }

//Delete node p, and then rebalance the tree.
    private void deleteEntry(TreeMapEntry<K,V> p) {
        modCount++;
        size--;

        // If strictly internal, copy successor's element to p and then make p
        // point to successor.
        if (p.left != null && p.right != null) {
            TreeMapEntry<K,V> s = successor(p);
            p.key = s.key;
            p.value = s.value;
            p = s;
        } // p has 2 children

        // Start fixup at replacement node, if it exists.
        TreeMapEntry<K,V> replacement = (p.left != null ? p.left : p.right);

        if (replacement != null) { //取p的一个子节点，为replacement 
            // Link replacement to parent
            replacement.parent = p.parent; //直接替换子节点replacement 的父节点为p的父节点
            if (p.parent == null) //如果p是root节点，root指向子节点replacement 
                root = replacement;
            else if (p == p.parent.left) //如果p是父节点的左节点，则父节点的左节点指向子节点replacement 
                p.parent.left  = replacement;
            else //否则p是父节点的右节点，则父节点的右节点指向子节点replacement 
                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) { // p没有子节点，且是唯一节点
            root = null;
        } else { //p是叶子节点，且不是root节点
            if (p.color == BLACK)
                fixAfterDeletion(p);

            if (p.parent != null) { 
                if (p == p.parent.left)//如果p是父节点的左节点，则父节点的左节点指向null
                    p.parent.left = null;
                else if (p == p.parent.right)//否则p是父节点的右节点，则父节点的右节点指向null
                    p.parent.right = null;
                p.parent = null;
            }
        }
    }