【Java 集合框架】TreeMap

68 阅读11分钟

TreeMap 是 Java Collections Framework 中的一员,是一个基于红黑树的 NavigableMap 实现。

TreeMap 根据不同的构造函数会对映射进行排序,默认是其键的自然顺序;在创建时如果提供了 Comparator ,则会根据 Comparator 进行排序。

这里说的映射是 Map 结构中的元素会映射到另一个数据结构中,这个数据结构负责处理排序,其中的元素称之为映射。在 TreeMap 中,映射数据结构是树结构。

该实现为 containsKey、get、put 和 remove 操作提供了保证的 log(n) 的时间复杂度。算法改编自 Cormen、Leiserson 和 Rivest 的《算法导论》。

注意,如果要正确实现 Map 接口,由树映射维护的排序(与任何排序映射一样)以及是否提供显式比较器都必须与 equals 一致。

这是因为 Map 接口是根据 equals 操作定义的,但排序映射使用其 compareTo (或 compare )方法执行所有键比较,因此,从排序映射的角度来看,该方法认为相等的两个键是相等的。

TreeMap 是无法保证同步的。如果多个线程并发地访问一个映射,并且其中至少有一个线程从结构上修改了映射,那么它必须从外部处理同步。(结构修改是添加或删除一个或多个映射的任何操作; 仅仅更改与现有键相关联的值并不是结构修改)。 这通常是通过在自然封装映射的某些对象上同步来实现的。如果不存在这样的对象,map 则应该使用Collections.synchronizedSortedMap 方法进行封装。

SortedMap m = Collections.synchronizedSortedMap(new TreeMap(...));

这个类的所有“集合视图方法”返回的集合的迭代器方法返回的迭代器是 fail-fast 的:如果映射在迭代器创建后的任何时间被结构修改,除了通过迭代器自己的remove方法以外,迭代器将抛出ConcurrentModificationException。

因此,在面对并发修改时,迭代器会快速而干净地失败,而不是在未来不确定的时间发生任意的、不确定的行为。

注意,迭代器的快速失败行为不能得到保证,因为一般来说,在存在不同步的并发修改时,不可能做出任何硬保证。

快速失败迭代器尽可能地抛出 ConcurrentModificationException。因此,编写一个依赖于此异常的程序是错误的:迭代器的快速失败行为应该只用于检测错误。

这个类及其视图中的方法返回的所有 Map.Entry 对都表示在产生映射时的快照。他们不支持 Entry.setValue 方法。(但是请注意,可以使用put更改关联映射中的映射)。

继承关系

public class TreeMap<K,V>
    extends AbstractMap<K,V>
    implements NavigableMap<K,V>, Cloneable, java.io.Serializable

TreeMap 继承自 AbstractMap ,实现了 NavigableMap、Cloneable 和 Serializable 接口。Cloneable 和 Serializable 很简单不用单独说明。

AbstractMap

AbstractMap 中对 Map 接口中定义的一些方法做了默认实现。无需展开说明。

NavigableMap

了解 TreeMap 之前,先看看他实现的 NavigableMap 接口中定义了什么能力。

public interface NavigableMap<K,V> extends SortedMap<K,V> {
    Map.Entry<K,V> lowerEntry(K key);
    K lowerKey(K key);
    Map.Entry<K,V> floorEntry(K key);
    K floorKey(K key);
    Map.Entry<K,V> ceilingEntry(K key);
    K ceilingKey(K key);
    Map.Entry<K,V> higherEntry(K key);
    K higherKey(K key);
    Map.Entry<K,V> firstEntry();
    Map.Entry<K,V> lastEntry();
    Map.Entry<K,V> pollFirstEntry();
    Map.Entry<K,V> pollLastEntry();
    NavigableMap<K,V> descendingMap();
    NavigableSet<K> navigableKeySet();
    NavigableSet<K> descendingKeySet();
    NavigableMap<K,V> subMap(K fromKey, boolean fromInclusive, K toKey,   boolean toInclusive);
    NavigableMap<K,V> headMap(K toKey, boolean inclusive);
    NavigableMap<K,V> tailMap(K fromKey, boolean inclusive);
    SortedMap<K,V> subMap(K fromKey, K toKey);
    SortedMap<K,V> headMap(K toKey);
    SortedMap<K,V> tailMap(K fromKey);
}

NavigableMap 继承自 SortedMap 接口,后者扩展了导航方法,为给定的搜索目标返回最接近的匹配。

可以将 NavigableMap 中定义的方法按组划分。

方法 lowerEntry、floorrentry、ceilingEntry 和 higherEntry 返回与键相关联的 Map.Entry 对象,它们分别处理小于、小于或等于、大于或等于和大于给定键,如果没有这样的键,则返回 null 。

    Map.Entry<K,V> lowerEntry(K key);
    Map.Entry<K,V> floorEntry(K key);
    Map.Entry<K,V> ceilingEntry(K key);
    Map.Entry<K,V> higherEntry(K key);

类似地,方法 lowerKey、floorKey、ceilingKey 和 higherKey 只返回相关的键。

    K lowerKey(K key);
    K floorKey(K key);
    K ceilingKey(K key);
    K higherKey(K key);

所有这些方法都是为定位而设计的,而不是遍历条目。可以按升序或降序键顺序访问和遍历 NavigableMap 。

    NavigableMap<K,V> descendingMap();

descendingMap 方法返回一个映射视图,其中包含所有相关和定向方法的倒置定义。升序操作和视图的性能可能比降序操作和视图快。

方法 subMap(K, boolean, K, boolean)、headMap(K, boolean) 和 tailMap(K, boolean) 在接受描述下界和上界是包含还是排除的附加参数方面与类似命名的SortedMap方法不同。

    NavigableMap<K,V> subMap(K fromKey, boolean fromInclusive, K toKey,   boolean toInclusive);
    NavigableMap<K,V> headMap(K toKey, boolean inclusive);
    NavigableMap<K,V> tailMap(K fromKey, boolean inclusive);

任何 NavigableMap 的子映射都必须实现 NavigableMap 接口。

该接口还定义了方法 firstEntry、pollFirstEntry、lastEntry 和 pollLastEntry ,它们返回and/or 删除最小和最大的映射(如果存在映射),否则返回null。

    Map.Entry<K,V> firstEntry();
    Map.Entry<K,V> lastEntry();
    Map.Entry<K,V> pollFirstEntry();
    Map.Entry<K,V> pollLastEntry();

入口返回方法的实现期望返回 Map.Entry对,表示生成映射时的快照,因此通常不支持可选的 Entry.setValue 方法。但是请注意,可以使用put方法更改关联映射中的映射。

subMap(K, K)、headMap(K) 和 tailMap(K) 方法被指定为返回 SortedMap ,以允许对 SortedMap 的现有实现进行兼容的改进以实现 NavigableMap,但鼓励该接口的扩展和实现覆盖这些方法以返回 NavigableMap 。

		// 由 SortedMap 中定义,子类需实现
    SortedMap<K,V> subMap(K fromKey, K toKey);
    SortedMap<K,V> headMap(K toKey);
    SortedMap<K,V> tailMap(K fromKey);

类似地,可以重写 keySet()以返回 NavigableSet 。

    NavigableSet<K> navigableKeySet();
    NavigableSet<K> descendingKeySet();

属性

    /**
     * 用于维持树映射中的顺序的比较器,如果使用键的自然顺序,则为空。
     */
    private final Comparator<? super K> comparator;

    private transient TreeMapEntry<K,V> root;

    /**
     * 树中的 entry 数量
     */
    private transient int size = 0;

    /**
     * 树的结构修改的计数。
     */
    private transient int modCount = 0;

通过属性,可以快速了解 TreeMap 的一些特性:

  • comparator 比较器表示支持排序。
  • root 代表底层的映射数据结构。
  • size 代表元素数量。
  • modCount 意味着可能会存在一些同步问题。

底层数据结构

TreeMapEntry 类代表当前树中的节点。

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;

    TreeMapEntry(K key, V value, TreeMapEntry<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 TreeMap(Comparator<? super K> comparator) {
        this.comparator = comparator;
    }

    // 给定 map 创建 treeMap ,key 自然排序
    public TreeMap(Map<? extends K, ? extends V> m) {
        comparator = null;
        putAll(m);
    }
    // 给定 SortedMap 创建 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) {
        }
    }

TreeMap 的构造方法主要是针对排序的区别。

核心逻辑

添加

添加新元素通过 put(key, value) 方法进行:

public V put(K key, V value) {
    TreeMapEntry<K,V> t = root;
    if (t == null) {
        compare(key, key); // 类型(可能是null)检查

        root = new TreeMapEntry<>(key, value, null);
        size = 1;
        modCount++;
        return null;
    }
    int cmp;
    TreeMapEntry<K,V> parent;
    // split comparator and comparable paths
    Comparator<? super K> 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();
        @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);
    }
    TreeMapEntry<K,V> e = new TreeMapEntry<>(key, value, parent);
    if (cmp < 0)
        parent.left = e;
    else
        parent.right = e;
    fixAfterInsertion(e);
    size++;
    modCount++;
    return null;
}

put 方法的返回值为先前 key 关联的值。

这个方法可以划分为几个步骤:

  1. 检查 root 是否存在,不存在则将新添加的键值对作为 root 节点,并返回 null (先前不存在值):

    		TreeMapEntry<K,V> t = root;
        if (t == null) {
            compare(key, key); // 类型(可能是null)检查
            root = new TreeMapEntry<>(key, value, null);
            size = 1;
            modCount++;
            return null;
        }
    

    这里的 compare(key, key) 的作用是使用 TreeMap 的正确比较方法来比较两个键 :

        final int compare(Object k1, Object k2) {
            return comparator==null ? ((Comparable<? super K>)k1).compareTo((K)k2)
                : comparator.compare((K)k1, (K)k2);
        }
    
  2. TreeMap 存在比较器,通过比较器进行遍历节点进行查找 :

        int cmp;
        TreeMapEntry<K,V> parent;
        // split comparator and comparable paths
        Comparator<? super K> 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);
        }
    

    cmp 代表查找结果,小于 0 代表 key 的顺序在当前节点的左子树中,大于 0 表示 key 的顺序在当前节点的右子树中。等于 0 表示查找到了 key 。

    所以查找到了直接 setValue 并返回之前的值。

  3. 不存在比较器的情况下,通过 key 进行查找,原理和比较器一样,只不过比较逻辑不同:

        if (cpr != null) {
            // ...
        } else {
            if (key == null)
                throw new NullPointerException();
            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);
        }
    
  4. 如果到这一步还没有 return ,代表当前树中不存在键为 key 的元素。那么就创建新的节点对象,并将其插入到树结构中:

    		TreeMapEntry<K,V> e = new TreeMapEntry<>(key, value, parent);
        if (cmp < 0)
            parent.left = e;
        else
            parent.right = e;
        fixAfterInsertion(e);
        size++;
        modCount++;
        return null;
    

    先构造元素对象,然后根据比较结果 cmp 判断是将其插入到父节点的左子节点还是右子节点。

    然后调用 fixAfterInsertion(e) 对红黑树进行重新平衡(旋转、染色)处理。

查询

查询元素主要是通过 get(key) 方法,通过 Key 查找键值对的值。

    public V get(Object key) {
        TreeMapEntry<K,V> p = getEntry(key);
        return (p==null ? null : p.value);
    }
final TreeMapEntry<K,V> getEntry(Object key) {
    // 卸载基于比较器的版本,以提高性能
    if (comparator != null)
        return getEntryUsingComparator(key);
    if (key == null)
        throw new NullPointerException();
    @SuppressWarnings("unchecked")
        Comparable<? super K> k = (Comparable<? super K>) key;
    TreeMapEntry<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;
}

getEntry 方法中,如果比较器 comparator 不为空,通过 getEntryUsingComparator 方法进行查询;否则通过 key 的自然排序顺序进行查询。从 root 节点开始进行查找,因为红黑树是二叉排序树,所以查询时间复杂度为 log(n) 。

Key 自然排序查找算法

核心的查找逻辑是 while 循环:

    TreeMapEntry<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;
    }

存在比较器的查找

存在比较器的情况下,通过 getEntryUsingComparator 方法进行查找:

final TreeMapEntry<K,V> getEntryUsingComparator(Object key) {
  	K k = (K) key;
    Comparator<? super K> cpr = comparator;
    if (cpr != null) {
        TreeMapEntry<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;
}

查找算法也是通过 compare 进行比较,只不过这里是通过比较器的 compare 方法进行的。

删除

    public V remove(Object key) {
        TreeMapEntry<K,V> p = getEntry(key);
        if (p == null)
            return null;

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

删除方法首先进行查找,如果找到通过 deleteEntry 方法进行删除:

    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) {
            // 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);
        } else if (p.parent == null) { // return if we are the only node.
            root = 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;
            }
        }
    }

删除方法也可以分为几部分进行:

  1. 首先是计数器更新:

        modCount++;
        size--;
    
  2. 待删除节点 p 存在左右子节点,通过 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
    
  3. 如果存在可替换的节点,则在替换节点启动 fixup。

        if (replacement != null) {
            // 进行替换
            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;
    
            //清空 p 的引用,以便 p 可以通过 fixAfterDeletion 删除
            p.left = p.right = p.parent = null;
    
            // Fix replacement
            if (p.color == BLACK)
                fixAfterDeletion(replacement);
        }
    
  4. replacement 不存在,但 p.parent 为空时,代表当前只存在一个节点,直接将 root 置为 null :

        if (replacement != null) {
            // ...
        } else if (p.parent == null) { // return if we are the only node.
            root = null;
        }
    
  5. 其他情况,没有子节点,如果当前时黑叶子,直接删除。并将叶子节点的父节点的子节点引用进行删除。

        if (replacement != null) {
            // ...
        } else if (p.parent == 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;
            }
        }
    

可以看出,在删除时,根据节点的类型不同,划分为三种情况:

  • 待删除的节点是根节点的情况
  • 待删除的节点既不是根也不是叶子的情况
  • 待删除的节点是叶子节点的情况

针对不同情况的特征分类处理是树结构处理逻辑中比较基础并且重要的思路。

排序

在处理删除时,通过 successor 方法获取了待删除节点的下一项。很明显,TreeMap 中存在一些处理元素顺序的逻辑。TreeMap 中,对应排序的逻辑,实质上就是红黑树的排序逻辑。

successor 方法返回了后一个节点,predecessor 方法返回前一个节点。

    static <K,V> TreeMapEntry<K,V> successor(TreeMapEntry<K,V> t) {
      	
        if (t == null)
          	// // 目标元素为空,直接 return 
            return null;
        else if (t.right != null) {
          	// 目标元素的右子树不为空,开始遍历右子树,查找最左边的节点
            TreeMapEntry<K,V> p = t.right;
            while (p.left != null)
                p = p.left;
            return p;
        } else {
          	// 目标元素不存在右子树,也不为空时
            TreeMapEntry<K,V> p = t.parent;
            TreeMapEntry<K,V> ch = t;
          	// 父节点不为空且当前节点是父节点的右子树时,继续向上查找,直到当前节点是左子树时,返回父节点
            while (p != null && ch == p.right) {
                ch = p;
                p = p.parent;
            }
            return p;
        }
    }

predecessor 方法同理,查找方向不同:

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

排序这块可以去学习红黑树的原理。

数据结构

在 TreeMap 中,不仅仅是一个红黑树结构来保存数据的,结构视图还包括:

  • entrySet (节点元素集合)
  • navigableKeySet( Key 的升序集合)
  • descendingMap (降序 Map)
    private transient EntrySet entrySet;
    private transient KeySet<K> navigableKeySet;
    private transient NavigableMap<K,V> descendingMap;

		// key 的升序集合
    public Set<K> keySet() {
        return navigableKeySet();
    }

		//  key 的降序集合
    public NavigableSet<K> descendingKeySet() {
        return descendingMap().navigableKeySet();
    }
		
		// 取自 AbstractMap 中的 values 「transient Collection<V> values;」
    public Collection<V> values() {
        Collection<V> vs = values;
        if (vs == null) {
            vs = new Values();
            values = vs;
        }
        return vs;
    }

    public Set<Map.Entry<K,V>> entrySet() {
        EntrySet es = entrySet;
        return (es != null) ? es : (entrySet = new EntrySet());
    }

    public NavigableMap<K, V> descendingMap() {
        NavigableMap<K, V> km = descendingMap;
        return (km != null) ? km :
            (descendingMap = new DescendingSubMap<>(this,
                                                    true, null, true,
                                                    true, null, true));
    }

这里的 EntrySet 结构继承自 AbstractSet ,后者是 Set 接口的框架实现:

class EntrySet extends AbstractSet<Map.Entry<K,V>> 

KeySet 同样也是,但它实现了 NavigableSet ,具有了排序能力:

static final class KeySet<E> extends AbstractSet<E> implements NavigableSet<E>

总结

以上就是 TreeMap 的核心逻辑,TreeMap 本质上是红黑树的封装,时间复杂度为 log(n) ,不支持同步。

红黑树高效的特性十分适合用来做 Map 的底层数据结构。