TreeMap的整体架构

TreeMap 底层的数据结构是红黑树，和 HashMap 的红黑树结构一样。
不同的是，TreeMap 利用了红黑树左节点小，右节点大的性质，根据 key 进行排序，使每个元素能够插入到红黑树大小适当的位置，维护了 key 的大小关系，TreeMap 适用于 key 需要排序的场景。
因为底层的数据结构是红黑树，所以 containsKey()、get()、put()、remove() 等方法的时间复杂度都是 log(n)

TreeMap 的部分底层源码

public class TreeMap<K,V> extends AbstractMap<K,V>
    implements NavigableMap<K,V>, Cloneable, java.io.Serializable {
    // 用于在 TreeMap 中保持顺序的外部比较器，如果它使用键的自然顺序，则该值为 null
    // 如果该值为 null，则使用 key 实现的 Comparable.compareTo()
    private final Comparator<? super K> comparator;
    // 红黑树的根节点
    private transient Entry<K,V> root;
    // 红黑树节点的数量
    private transient int size = 0;
    // 树结构修改的次数
    private transient int modCount = 0;
    // 使用键的自然顺序构造一个新的空 TreeMap。
    // 键的类型必须实现 Comparable 接口
    // 所有的键相互比较 k1.compareTo(k2) 必须不抛出 ClassCastException
    public TreeMap() {
        comparator = null;
    }
    // 根据给定的外部排序器构造一个新的空 TreeMap
    // 所有键相互比较由给定的比较器 comparator.compare(k1, k2) 必须不抛出 ClassCastException
    public TreeMap(Comparator<? super K> comparator) {
        this.comparator = comparator;
    }
    // 构造一个包含给定 map 数据的 TreeMap
    // 根据键的自然顺序排序
    // 键的类型必须实现 Comparable 接口
    // 所有的键相互比较 k1.compareTo(k2) 必须不抛出 ClassCastException
    // 该方法运行时间为 n*log(n)
    public TreeMap(Map<? extends K, ? extends V> m) {
        comparator = null;
        putAll(m);
    }
    // 构造一个包含给定 map 数据的 TreeMap，并使用与指定的 SortedMap 相同的顺序
    // 该方法在线性时间内运行  
    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) {
        }
    }
    // Red-black mechanics
    private static final boolean RED   = false;
    private static final boolean BLACK = true;
    // 红黑树的节点
    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;
        Entry(K key, V value, Entry<K,V> parent) {
            this.key = key;
            this.value = value;
            this.parent = parent;
        }
        public K getKey() {
            return key;
        }
        public V getValue() {
            return value;
        }
        public V setValue(V value) {
            V oldValue = this.value;
            this.value = value;
            return oldValue;
        }
        public boolean equals(Object o) {
            if (!(o instanceof Map.Entry))
                return false;
            Map.Entry<?,?> e = (Map.Entry<?,?>)o;
            return valEquals(key,e.getKey()) && valEquals(value,e.getValue());
        }
        public int hashCode() {
            int keyHash = (key==null ? 0 : key.hashCode());
            int valueHash = (value==null ? 0 : value.hashCode());
            return keyHash ^ valueHash;
        }
        public String toString() {
            return key + "=" + value;
        }
    }

TreeMap的新增

TreeMap 新增节点的大概步骤：

判断红黑树的跟节点是否为空，为空的话，新增的节点直接作为根节点
根据红黑树左小右大的特性，进行判断，找到应该新增节点的父节点，或者直接替换旧值
在父节点的左边或右边插入新增节点
着色旋转，达到平衡，结束

put() 的底层源码实现如下：

// 将指定值与此 TreeMap 中的指定键关联。如果之前的 TreeMap 包含键的映射，那么旧的值将被替换
// 返回 key 之前映射的值，如果不存在则返回 null
public V put(K key, V value) {
    // 红黑树的根节点
    Entry<K,V> t = root;
    // 判断红黑树的跟节点是否为空，如果为 null，则新增节点直接作为根节点
    if (t == null) {
        // 使用该 TreeMap 的正确比较方法比较两个键
        // compare 方法限制了 key 不能为 null(但 val 可以为 null)
        compare(key, key); // type (and possibly null) check
        root = new Entry<>(key, value, null);
        size = 1;
        modCount++;
        return null;
    }
    //cmp 用于存储每一次节点的 key 对比的结果
    int cmp;
    // 应该新增节点的父节点
    Entry<K,V> parent;
    // split comparator and comparable paths
    Comparator<? super K> cpr = comparator;
    // 如果传入的外部排序器 cpmparator 不为 null，则使用外部排序器
    if (cpr != null) {
        do {
            // 第一次循环时，parent 是红黑树的根节点，一次循环结束时，parent 就是上次比过的对象
            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;
            // 如果新增节点的 key == 当前节点的 key，不用新增节点，直接替换旧值
            else
                return t.setValue(value);
        } while (t != null);
    }
    // 如果传入的外部排序器 cpmparator 为 null，则使用 key 实现的 Comparable.compareTo()
    else {
        if (key == null)
            throw new NullPointerException();
        @SuppressWarnings("unchecked")
        Comparable<? super K> k = (Comparable<? super K>) key;
        // 这个循环和上面 if 中的循环类似，不做注释说明
        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);
    // 此时 cmp 代表最后一次对比的结果
    // 如果 cmp < 0 ，代表节点 e 的 key < 节点 parent 的 key
    if (cmp < 0)
        // 节点 e 作为节点 parent 的左子节点
        parent.left = e;
    else
        parent.right = e;
    // 进行着色旋转操作
    fixAfterInsertion(e);
    // 红黑树的节点数量 + 1
    size++;
    // 红黑树的结构的修改次数 + 1
    modCount++;
    return null;
}
// 使用该 TreeMap 的正确比较方法比较两个键
final int compare(Object k1, Object k2) {
    return comparator==null ? ((Comparable<? super K>)k1).compareTo((K)k2)
        : comparator.compare((K)k1, (K)k2);
}
/** From CLR */
// 该方法进行着色旋转操作
private void fixAfterInsertion(Entry<K,V> x) {
    x.color = RED;
    while (x != null && x != root && x.parent.color == RED) {
        if (parentOf(x) == leftOf(parentOf(parentOf(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;
}