Java编程拾遗『TreeMap』

通过之前两节的讲述，相信对二叉树、以及红黑树有了初步的认识。本篇文章来讲述一下基于红黑树实现的Map—TreeMap。之前介绍HashMap时，我们都知道，HashMap键值对之间没有特定的顺序。结合之前讲的红黑树特点，我们也可以猜测到，TreeMap是可以保证键值对之间的顺序的。

TreeMap继承了AbstractMap抽象类，实现了NavigableMap、Cloneable、Serializable接口。对于抽象类AbstractMap，Clobeable、Serializable接口的作用，比较简单，这里不多讲了。简单提一下两个特殊的接口SortedMap、NavigableMap，这两个接口主要是根据TreeMap有序的特点，获取key前后或一定范围内的键值集合或键值对集合，接下来会详细介绍。

1. 使用规则

1.1 TreeMap基本定义

public class TreeMap<K,V>
    extends AbstractMap<K,V>
    implements NavigableMap<K,V>, Cloneable, java.io.Serializable
{
    private final Comparator<? super K> comparator;

    private transient Entry<K,V> root;

    private transient int size = 0;

    private transient int modCount = 0;

    //methods

}    

继承关系上面已经讲过了，这里看一下TreeMap的主要成员变量。

• comparator，用于key排序的比较器
• root，用于表示红黑树结构的红黑树根节点
• size， TreeMap中键值对的数目
• modCount：TreeMap的修改次数，用于实现fail-fast机制，抛出ConCurrentModificationException

1.2 构造函数

1.3 SortedMap

TreeMap实现了NavigableMap接口，NavigableMap接口又实现了SortedMap接口，所以TreeMap实现的NavigableMap接口的方法，有一部分是来自SortedMap的，下面来看一下SortedMap接口的方法。首先看一下接口定义：

public interface SortedMap<K,V> extends Map<K,V> {
	Comparator<? super K> comparator();

	SortedMap<K,V> subMap(K fromKey, K toKey);

	SortedMap<K,V> headMap(K toKey);

	SortedMap<K,V> tailMap(K fromKey);

	K firstKey();

	K lastKey();

	Set<K> keySet();

	Collection<V> values();

	Set<Map.Entry<K, V>> entrySet();
}

1.3 NavigableMap

1.3 Map

2. 实现源码分析

TreeMap内部是用红黑树实现的，关于红黑树的相关特性以及Java实现上一篇文章已经详细讲述过。本篇文章结合上篇文章讲述的红黑树的特性，来看一下TreeMap的实现。

2.1 构造函数

public TreeMap() {
    comparator = null;
}

默认构造函数，对TreeMap内部成员变量comparator赋null，该情形下需要TreeMap中key对应的类实现Comparable接口。

public TreeMap(Comparator<? super K> comparator) {
    this.comparator = comparator;
}

通过Comparator对象初始化TreeMap内部成员变量comparator，TreeMap内部成员变量comparator用于确定键值对之间的顺序。

public TreeMap(Map<? extends K, ? extends V> m) {
    comparator = null;
    putAll(m);
}

通过Map初始化TreeMap，要求Map的key实现了Comparable接口。对于putAll方法，下面单独讲述。

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

通过SortedMap初始化TreeMap，buildFromSorted函数，会通过一个有序序列构建一个红黑树。

2.2 红黑树节点Entry

关于红黑树的Java实现，上篇文章已经将结果，其实Java API中的实现跟我上节列出来的实现差不多。在TreeMap内部，使用内部类Entry来表示红黑树的节点，如下：

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;

    /**
     * Make a new cell with given key, value, and parent, and with
     * {@code null} child links, and BLACK color.
     */
    Entry(K key, V value, Entry<K,V> parent) {
        this.key = key;
        this.value = value;
        this.parent = parent;
    }

    /**
     * Returns the key.
     *
     * @return the key
     */
    public K getKey() {
        return key;
    }

    /**
     * Returns the value associated with the key.
     *
     * @return the value associated with the key
     */
    public V getValue() {
        return value;
    }

    /**
     * Replaces the value currently associated with the key with the given
     * value.
     *
     * @return the value associated with the key before this method was
     *         called
     */
    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;
    }
}

每个节点除了键(key)和值(value)之外，还有三个引用，分别指向其左孩子(left)、右孩子(right)和父节点(parent)，对于根节点，父节点为null，对于叶子节点，孩子节点都为null，还有一个成员color表示颜色，TreeMap是用红黑树实现的，每个节点都有一个颜色，非黑即红。

2.3 put(K key, V value)

public V put(K key, V value) {
    Entry<K,V> t = root;
    if (t == null) {
    	//红黑树为空，直接初始化一个Entry节点作为跟节点即可
        compare(key, key); // type (and possibly null) check

        root = new Entry<>(key, value, null);
        size = 1;
        modCount++;
        return null;
    }
    int cmp;
    Entry<K,V> parent;
    // split comparator and comparable paths
    Comparator<? super K> cpr = comparator;
    if (cpr != null) {
    	//TreeMap内部comparator不为null，遍历红黑树，查看是否存在key节点，存在则替换value
        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 {
    	//TreeMap内部comparator为null，遍历红黑树，查看是否存在key节点，存在则替换value
        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);
    }
    //红黑树中不存在key节点，则在红黑树添加一个节点
    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;
}

上述两种遍历红黑树的目的是一致的，就是如果在遍历过程中找到key节点，则进行value替换，如果没找到，则最终定位到要插入节点的父节点parent。

2.4 putAll(Map<? extends K, ? extends V> map)

public void putAll(Map<? extends K, ? extends V> map) {
    int mapSize = map.size();
    if (size==0 && mapSize!=0 && map instanceof SortedMap) {
        Comparator<?> c = ((SortedMap<?,?>)map).comparator();
        if (c == comparator || (c != null && c.equals(comparator))) {
            //如果TreeMap为空，且待插入的Map为SortedMap，可以直接调用buildFromSorted方法构建红黑树
            ++modCount;
            try {
                buildFromSorted(mapSize, map.entrySet().iterator(),
                                null, null);
            } catch (java.io.IOException cannotHappen) {
            } catch (ClassNotFoundException cannotHappen) {
            }
            return;
        }
    }
    //调用AbstractMap的putAll方法，依次插入map中的元素，问题在于每次插入都需要调整红黑树结构，效率不高
    super.putAll(map);
}

putAll方法用于将一个Map的所有元素添加到TreeMap中，如果TreeMap为空，且map是SortedMap，参数map的比较器和当前TreeMap的比较器相同，那么使用迭代器从小到大以折半的方式将树组装成一个新的红黑树，通过buildFromSorted方法构建的红黑树，最下方的节点是红色，其他的都是黑色。否则会调用AbstractMap的putAll方法，依次插入map中的元素，每一次插入都需要修复红黑树结构，使之符合红黑树约束，大致保持平衡，效率不高。在TreeMap的构造函数中，通过Map初始化TreeMap就是调用的putAll方法，如果Map符合上述约束，肯定会调用buildFromSorted方法构建红黑树，否则将调用putAll方法依次插入。下面来看一下buildFromSorted方法实现：

/**
 * @param size 从iterator或stream中读取的key或key-value对的数目
 * @param it 如果非null, 则从iterator获取key-value entries或keys插入TreeMap
 * @param str 如果非null, 则从ObjectInputStream中读取keys插入TreeMap(defaultVal为null时，value的值默认跟key相等，取自stream)
 *        必须严格保证it和str其中一个不为null
 * @param defaultVal 如果非null, 则插入TreeMap中的键值对value都为defaultVal
 *        如果为null, 则从iterator或stream中读取value
 */
private void buildFromSorted(int size, Iterator<?> it,
                             java.io.ObjectInputStream str,
                             V defaultVal)
    throws  java.io.IOException, ClassNotFoundException {
    this.size = size;
    root = buildFromSorted(0, 0, size-1, computeRedLevel(size),
                           it, str, defaultVal);
}

//计算红黑树叶子结点的层数，在递归过程中，如果level == redLevel，则此时的结点为红色
private static int computeRedLevel(int sz) {
    int level = 0;
    for (int m = sz - 1; m >= 0; m = m / 2 - 1)
        level++;
    return level;
}

/**
 * @param level 当前数的level，初始值为0
 * @param lo 当前子树第一个元素的index，初始值为0
 * @param hi 当前子树最后一个元素的index，初始值为size-1
 * @param redLevel 红黑树树节点为红的level
 */
@SuppressWarnings("unchecked")
private final Entry<K,V> buildFromSorted(int level, int lo, int hi,
                                         int redLevel,
                                         Iterator<?> it,
                                         java.io.ObjectInputStream str,
                                         V defaultVal)
    throws  java.io.IOException, ClassNotFoundException {

    if (hi < lo) return null;

    int mid = (lo + hi) >>> 1;

    Entry<K,V> left  = null;
    if (lo < mid)
        left = buildFromSorted(level+1, lo, mid - 1, redLevel,
                               it, str, defaultVal);

    // extract key and/or value from iterator or stream
    K key;
    V value;
    if (it != null) {
        if (defaultVal==null) {
            Map.Entry<?,?> entry = (Map.Entry<?,?>)it.next();
            key = (K)entry.getKey();
            value = (V)entry.getValue();
        } else {
            key = (K)it.next();
            value = defaultVal;
        }
    } else { // use stream
        key = (K) str.readObject();
        value = (defaultVal != null ? defaultVal : (V) str.readObject());
    }

    Entry<K,V> middle =  new Entry<>(key, value, null);

    // color nodes in non-full bottommost level red
    if (level == redLevel)
        middle.color = RED;

    if (left != null) {
        middle.left = left;
        left.parent = middle;
    }

    if (mid < hi) {
        Entry<K,V> right = buildFromSorted(level+1, mid+1, hi, redLevel,
                                           it, str, defaultVal);
        middle.right = right;
        right.parent = middle;
    }

    return middle;
}

该方法其实是通过一个有序序列构建红黑树，其实就是通过递归的思想构建树。比如存在有序序列[1,2,3,4,5,6,7,8 ,9,10]，通过上述方法构建红黑树的过程为：

1. 以最中间的数作为根结点，然后将序列分成两组，(1,2,3,4) (6,7,8,9,10)
2. 递归在第一组序列中找出最中间的树作为根结点，建立一个子树，该子树作为整个树的左子树
3. 递归在第二组序列中找出最中间的树作为根结点，建立一个子树，该子树作为整个树的右子树
4. 最终形成的树是在叶子结点以上是一个满二叉树，把所有的叶子节点标记为红色，满足红黑树的大致平衡要求

上述方法中，lo和hi参数是从迭代器或stream中取出元素的最小和最大索引（元素用于插入当前子树）。实际上没有被索引，只是因为序列是有序的，可以通过这种折半取数的方式保证顺序。

如下图所示：

如果不满足map是SortedMap，当前的TreeMap为空，且参数map的比较器和当前TreeMap的比较器相同的条件，则需要调用AbstractMap的putAll方法，依次调用TreeMap的put方法插入元素，每次插入都需要调整红黑树结构，效率不高，时间复杂度为O(n * logn)，AbstractMap的putAll方法如下：

public void putAll(Map<? extends K, ? extends V> m) {
    for (Map.Entry<? extends K, ? extends V> e : m.entrySet())
        put(e.getKey(), e.getValue());
}

2.4 remove(Object key)

/**
 * 从TreeMap中删除key对应的键值对
 *
 * @param  key key 待删除键值对的key
 * @return 待删除key对应的value，如果key不存在，或对应的value为null，返回null
 * @throws ClassCastException 如果key和TreeMap中所有的key无法进行比较，抛ClassCastException
 * @throws NullPointerException 如果key为null并且TreeMapcompar为null，或者TreeMap的comparator不允许null key，抛NullPointerException
 */
public V remove(Object key) {
	//获取key对应的Entry节点
    Entry<K,V> p = getEntry(key);
    if (p == null)
        return null;

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

getEntry方法，用于获取key对应的Entry节点，如下：

/**
* 根据key获取key对应的Entry红黑树节点
*/
final Entry<K,V> getEntry(Object key) {
    if (comparator != null)
    	//comparator不为null，则调用getEntryUsingComparator方法获取节点
        return getEntryUsingComparator(key);
    if (key == null)
    	//comparator为null(按自然排序)，如果key为null，抛NPE
        throw new NullPointerException();

    @SuppressWarnings("unchecked")
    Comparable<? super K> k = (Comparable<? super K>) key;
    Entry<K,V> p = root;
    //遍历红黑树，获取Entry节点
    while (p != null) {
        int cmp = k.compareTo(p.key);
        if (cmp < 0)
            //cmp<0，表示key小于当前Entry节点的key，所以遍历当前节点的左子树
            p = p.left;
        else if (cmp > 0)
            //cmp>0，表示key大于当前Entry节点的key，所以遍历当前节点的右子树
            p = p.right;
        else
            //cmp == 0，表示key与当前Entry检点的key相等，直接返回当前Entry节点p即可
            return p;
    }
    return null;
}

/**
* TreeMap的comparator非null时，遍历红黑树获取key对应的Entry红黑树节点
*/
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;
        //遍历红黑树，获取Entry节点
        while (p != null) {
            int cmp = cpr.compare(k, p.key);
            if (cmp < 0)
            	//cmp<0，表示key小于当前Entry节点的key，所以遍历当前节点的左子树
                p = p.left;
            else if (cmp > 0)
            	//cmp>0，表示key大于当前Entry节点的key，所以遍历当前节点的右子树
                p = p.right;
            else
            	//cmp == 0，表示key与当前Entry检点的key相等，直接返回当前Entry节点p即可
                return p;
        }
    }
    return null;
}

下面看一下deleteEntry方法的实现，删除的算法，节点有三种情况：

• 叶子节点：这个容易处理，直接修改父节点对应引用置null即可。
• 只有一个孩子：就是在父亲节点和孩子节点直接建立链接。
• 有两个孩子：先找到后继，找到后，替换当前节点的内容为后继节点，然后再删除后继节点，因为这个后继节点一定没有左孩子，所以就将两个孩子的情况转换为了前面两种情况。

private void deleteEntry(Entry<K,V> p) {
    modCount++;
    size--;

    /*处理节点有两个孩子的情况，s为后继，当前节点p的key和value设置为了s的key和value，
    然后将待删节点p指向了s，这样就转换为了一个孩子或叶子节点的情况。*/
    if (p.left != null && p.right != null) {
        Entry<K,V> s = successor(p);
        p.key = s.key;
        p.value = s.value;
        p = s;
    }

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

    /*p为待删节点，replacement为要替换p的孩子节点，主体代码就是在p的父节点p.parent和replacement之间建立链接，
    以替换p.parent和p原来的链接，如果p.parent为null，则修改root以指向新的根。fixAfterDeletion重新平衡树。*/
    if (replacement != null) {
        // Link replacement to parent
        replacement.parent = p.parent;
        if (p.parent == null)
            root = replacement;
        else if (p == p.parent.left)
            p.parent.left  = replacement;
        else
            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);
    } 
    /*再具体分为两种情况，一种是删除最后一个节点，则直接修改root为null*/
    else if (p.parent == null) { // return if we are the only node.
        root = null;
    }
    /*否则就是根据待删节点是父节点的左孩子还是右孩子，相应的设置孩子节点为null*/
    else { //  No children. Use self as phantom replacement and unlink.
       if (p.color == BLACK)
           fixAfterDeletion(p);

       if (p.parent != null) {
           if (p == p.parent.left)
               p.parent.left = null;
           else if (p == p.parent.right)
               p.parent.right = null;
           p.parent = null;
       }
    }
}

2.5 get(Object key)

根据key，获取key对应的value，如果key不存在或者对应的value为null，返回null。如下：

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

通过getEntry方法获取key对应的Entry，如果Entry不为null，则说明TreeMap中存在key对应的键值对，直接返回返回Entry的value即可。getEntry方法在上面讲remove方法时，已经讲过了，这里不重复讲述了。

2.6 containsKey(Object key)

判断当前TreeMap中是否存在键值为key的键值对，如下：

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

通过getEntry方法获取key对应的Entry，如果Entry不为null，则说明TreeMap中存在key对应的键值对。

2.7 containsValue(Object value)

判断当前TreeMap中是否存在值为value的键值对，如下：

public boolean containsValue(Object value) {
    for (Entry<K,V> e = getFirstEntry(); e != null; e = successor(e))
        if (valEquals(value, e.value))
            return true;
    return false;
}

getFirstEntry方法用于获取当前TreeMap中key最小的Entry，如下：

final Entry<K,V> getFirstEntry() {
    Entry<K,V> p = root;
    if (p != null)
        while (p.left != null)
            p = p.left;
    return p;
}

从跟节点一直往左遍历，到最后肯定就是最小的Entry。

successor方法用于获取当前节点的后继几点，如下：

static <K,V> TreeMap.Entry<K,V> successor(Entry<K,V> t) {
    if (t == null)
        return null;
    else if (t.right != null) {
        Entry<K,V> p = t.right;
        while (p.left != null)
            p = p.left;
        return p;
    } else {
        Entry<K,V> p = t.parent;
        Entry<K,V> ch = t;
        while (p != null && ch == p.right) {
            ch = p;
            p = p.parent;
        }
        return p;
    }
}

算法描述为：

• 如果有右孩子(t.right!=null)，则后继为右子树中最小的节点。
• 如果没有右孩子，后继为某祖先节点，从当前节点往上找，如果它是父节点的右孩子，则继续找父节点，直到它不是右孩子或父节点为空，第一个非右孩子节点的父亲节点就是后继节点，如果父节点为空，则后继为null。

如下图所示：

 3. TreeMap特点

TreeMap继承了AbstractMap抽象类(相当于间接实现了Map接口)，除此之外还实现了NavigableMap接口。所以相比于HashMap，TreeMap除了存储了键值对之外，还可以保持键值对之间的先后顺序。内部使用红黑树实现，效率较高。

• TreeMap可以保证内部key的顺序，可以很方便的获取第一个、最后一个、或某区间内的Entry/key
• 根据键查找，删除，插入效率较高，时间复杂度为O(logn)

参考链接：

1. Java API源码
2. 《Java编程的逻辑》
3. 集合-TreeMap解析