集合包系列八 —— TreeMap

本系列文章所描述的所有类或接口都是基于 JDK 1.7的源码，在其它 JDK 的实现方式中可能会有所不同。

一、实现方式

TreeMap 是一个支持排序的 Map 实现，其实现方式和 HashMap 并不相同。

二、创建 TreeMap

在此步 TreeMap 只是将 comparator 属性赋值为 null，如希望控制 TreeMap 中元素的存储顺序，可使用带 Comparator 参数的构造器。

    public TreeMap() {
        comparator = null;
    }

    public TreeMap(Comparator comparator) {
        this.comparator = comparator;
    }

    public TreeMap(Map m) {
        comparator = null;
        putAll(m);
    }

    public TreeMap(SortedMap m) {
        comparator = m.comparator();
        try {
            buildFromSorted(m.size(), m.entrySet().iterator(), null, null);
        } catch (java.io.IOException cannotHappen) {
        } catch (ClassNotFoundException cannotHappen) {
        }
    }

三、添加元素 put(Object key, Object value)

当调用 put 时，先要判断 root 属性是否为 null，如是，则创建一个新的 Entry 对象，并赋值给 root 属性。如不是，则首先判断是否传入了指定的 Comparator 实现，如已传入，则基于红黑树的方式遍历，基于 comparator 来比较 key 应放在树的左边还是右边，如找到相等的 key，则直接替换其 value，并返回结束 put 操作，如没有找到相等的 key，则一直寻找到左边或右边节点为 null 的元素，如 comparator 实现为 null，则判断 key 是否为 null，是则抛出 NullPointerException，否则将 key 转化为 Comparable，进行与上面同样的遍历和比较过程。
　　通过上面的步骤，如未找到相同的 key，则进入以下过程，即创建一个新的 Entry 对象，并将其 parent 设置为上面所寻找到的元素，并根据和 parent key 比较的情况来设置 parent 的 left 或 right 属性。
　　综上所述，TreeMap 是一个典型的基于红黑树的实现，因此它要求一定要有 key 比较的方法，要么传入 Comparator 实现，要么 key 对象实现 Comparator 接口。

    public V put(K key, V value) {
        Entry t = root;
        if (t == null) {
            compare(key, key); // type (and possibly null) check

            root = new Entry<>(key, value, null);
            size = 1;
            modCount++;
            return null;
        }
        int cmp;
        Entry parent;
        // split comparator and comparable paths
        Comparator cpr = comparator;
        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);
        }
        else {
            if (key == null)
                throw new NullPointerException();
            Comparable k = (Comparable) 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 e = new Entry<>(key, value, parent);
        if (cmp < 0)
            parent.left = e;
        else
            parent.right = e;
        fixAfterInsertion(e);
        size++;
        modCount++;
        return null;
    }

四、获取元素 get(Object)

TreeMap 在查找 key 时就是个典型的红黑树查找过程，从根对象开始往下比较，一直找到相等的 key，并返回其 value。
　　和 put 时同样的处理方式，如未传入 Comparator 实现，当传入的 Object 为 null 时，则直接抛出 NullPointerException。

    public V get(Object key) {
        Entry p = getEntry(key);
        return (p==null ? null : p.value);
    }

    final Entry getEntry(Object key) {
        // Offload comparator-based version for sake of performance
        if (comparator != null)
            return getEntryUsingComparator(key);
        if (key == null)
            throw new NullPointerException();
        Comparable k = (Comparable) key;
        Entry 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;
    }

五、删除元素 remove(Object)

remove 首先要做的是 getEntry，然后则是将此 Entry 从红黑树上删除，并重新调整树上的相关的节点。

    public V remove(Object key) {
        Entry p = getEntry(key);
        if (p == null)
            return null;

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

六、判断元素是否存在 containsKey(Object)

和 get 方法一样，都通过 getEntry 方法来完成，因此过程和 get 方法一样，只是 containsKey 直接判断返回的 Entry 是否为 null，为 null 则返回 false，不为 null 则返回 true。

    public boolean containsKey(Object key) {
        return getEntry(key) != null;
    }

七、遍历元素 keySet()

调用 keySet 方法后返回 TreeMap 的内部类 KeySet 对象的实例，iterator 的遍历从根开始，基于红黑树顺序完成。

    public Set keySet() {
        return navigableKeySet();
    }

八、注意要点

对 TreeMap 而言，最应了解的有以下几点：

• TreeMap 基于红黑树实现，无容量限制；
• TreeMap 是非线程安全的。