TreeMap源码分析

TreeMap是NavbagableMap的实现，底层基于红黑树。这个Map按照Comparable将键值排序，或者按照在创建Map时提供的Compartor。
TreeMap的类继承关系图如下：

TreeMap与HashMap的一个重要区别是：TreeMap不支持键

Comparable与Comparator

Comparable接口用于自身与另外一个对象比较，其接口定义如下：

public interface Comparable<T> {
    public int compareTo(T o);
}

下面以Integer为例，该类实现了Comparable接口，其compare()方法如下：

public int compareTo(Integer anotherInteger) {
        return compare(this.value, anotherInteger.value);
    }

    //拿自己的Integer的值与另一个对象比较
    public static int compare(int x, int y) {
        return (x < y) ? -1 : ((x == y) ? 0 : 1);
    }

而Comparator接口表示比较两个对象，接口如下：

public interface Comparator<T> {
    int compare(T o1, T o2);

    ...//1.8增加的方法
}

Comparator一般用于作为Arrays或Collections排序时的一个参数，用于比较数组或集合中元素的顺序。
TreeMap是一个排序了的Map，所以依靠Comparator参数将元素排序。

构造方法

TreeMap中有一个compartor的字段如下：

  private final Comparator<? super K> comparator;

可以看到该字段为final，所以必须在构造方法中指定，如果为null，那么将使用键值的自然排序，否则使用该Compartor定下的规则。
所以TreeMap的每个构造方法中都对comparator字段进行了初始化，如下：

 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 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的底层是基于红黑树的实现，所以像get、put、remove、containsKey这些方法都会花费log(n)的时间复杂度。这儿不会着重于红黑树的具体实现以及转换，只要知道TreeMap的基本思路就可以了。

put操作

TreeMap的put方法如下所示：

public V put(K key, V value) {
        //得到红黑树根结点
        Entry<K,V> t = root;
        //如果Map为空
        if (t == null) {
            compare(key, key); // type (and possibly null) check

            //新建红黑树根结点
            root = new Entry<>(key, value, null);
            size = 1;
            modCount++;
            return null;
        }
        //如果Map不为空,找到插入新节点的父节点
        int cmp;
        Entry<K,V> parent;
        Comparator<? super K> cpr = comparator;
        //如果提供了Compartor
        if (cpr != null) {
            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);
        }
        //没有提供Comparator
        else {
            if (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;
    }

从put方法可以看到，有几步流程：
1. 如果Map为空，那么直接将新插入的值作为根结点。此时，如果提供了Compartor，就得看Compartor是否支持null键值；如果没有提供Compartor，那么将会抛出NullPointerException。
2. 如果Map不为空，那么需要找到新插入的键值的父节点。在查找过程中，如果遇到了键值相等的，那么将会调用Entry.setValue()更新值。
3. 一旦找到了父节点，那么插入新节点，尺寸+1

另外，我们需要注意到:
1. 如果没有提供Comparator，那么将不支持Null的键；如果提供了Comparator，那么是否支持Null的键将取决于Comparator的具体实现;
2. 如果没有提供Comparator，那么将会将键强制转换成Comparable接口，所以没有提供Comparator，那么键必须得实现Comparable接口，否则将抛出ClassCastException。

get操作

TreeMap的get方法根据键得到值，由于红黑树是一颗二叉搜索树，所以查询键值的操作很快，其实现如下：

 //根据键得到值，如果不存在，那么返回null
 public V get(Object key) {
        Entry<K,V> p = getEntry(key);
        return (p==null ? null : p.value);
    }

//根据键得到Entry节点，如果不存在，那么返回null
final Entry<K,V> getEntry(Object key) {
        //如果提供了Comparator
        if (comparator != null)
            return getEntryUsingComparator(key);
        //如果没有提供Comparator
        //由于put不允许键为null，那么get也不允许键为null
        if (key == null)
            throw new NullPointerException();

        @SuppressWarnings("unchecked")
            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;
    }

从上面可以看到，get()方法的流程：
1. 如果提供了Comparator，那么使用getEntryUsingComparator()方法
2. 如果没有提供Comparator，并且键为null，抛出NullPointerException
3. 如果没有提供Comparator且键不为null，将键强制转换为Comparable接口，如果键没有实现，那么抛出ClassCastExceotion
4. 如果没有提供Comparator且键不为null，且键实现了Comparable接口，那么从根结点开始遍历红黑树，一旦找到则返回节点，否则返回null

下面看一下getEntryUsingComparator()方法，该方法也是从根节点开始遍历，与使用Comparable接口的遍历类似：

final Entry<K,V> getEntryUsingComparator(Object key) {
        @SuppressWarnings("unchecked")
            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;
    }

remove()方法

remove()方法用于删除键对应的节点，如果不存在，则返回null，其内部实现包含了getEntry()查找节点的步骤，一旦找到了，再执行删除操作，实现如下：

 public V remove(Object key) {
        //查找节点
        Entry<K,V> p = getEntry(key);
        //如果Map中不包含节点，返回null
        if (p == null)
            return null;

        V oldValue = p.value;
        //删除节点
        deleteEntry(p);
        return oldValue;
    }

从代码可以看到，remove()方法主要有两步：
1. getEntry()得到待删除的节点
2. 如果Map中不包含节点，返回null；否则调用deleteEntry()删除节点

deleteEntry()方法中删除接节点后，为了维持红黑树的结构，还需要进行调整，这儿就不说明了。

总结

TreeMap中比较元素的规则可能来自于两方面，一方面是Comparator，另一方面是使用键值的Comparable接口中的方法。这两种方式有很大的区别：使用键的Comparable接口，那么就不允许null的键；而如果使用Comparator接口，那么是否支持null的键将由Comparator的实现决定。另外需要铭记的就是TreeMap的底层是使用的红黑树的结构，所以其get、put、remove方法中都涉及了树的操作，只要记住这一点，理解起各个方法时就容易多了。