深度解析TreeMap工作原理

一、引言

在 Java 的集合框架中，TreeMap 是一个独特且强大的 Map 实现。与 HashMap 不同，TreeMap 不仅能存储键值对，还能根据键的自然顺序或者指定的比较器顺序对键进行排序。这使得 TreeMap 在需要有序键值对存储和检索的场景中发挥着重要作用。本文将全方位、深入地解析 TreeMap 的原理，包括其底层数据结构、核心属性、构造方法、常用操作的实现细节等，并辅以丰富的代码示例。

二、TreeMap 概述

2.1 定义与用途

TreeMap 是 java.util 包下的一个类，实现了 NavigableMap 接口，而 NavigableMap 又继承自 SortedMap 接口。这意味着 TreeMap 中的键是有序的。它主要用于存储键值对，并且可以根据键的顺序对这些键值对进行高效的查找、插入和删除操作。常见的应用场景包括按时间顺序存储事件记录、按字母顺序存储字典条目等。

2.2 继承关系与实现接口

TreeMap 的继承关系如下：

java.lang.Object
    └─ java.util.AbstractMap<K,V>
        └─ java.util.TreeMap<K,V>

它实现了 NavigableMap<K,V>、Cloneable 和 java.io.Serializable 接口，具备克隆和序列化的能力，同时支持导航操作，如查找大于或小于某个键的键值对等。

import java.util.TreeMap;
import java.util.Map;

public class TreeMapOverview {
    public static void main(String[] args) {
        // 创建一个 TreeMap 对象
        TreeMap<String, Integer> treeMap = new TreeMap<>();
        // 可以将其赋值给 Map 接口类型的变量
        Map<String, Integer> map = treeMap;
    }
}

三、底层数据结构：红黑树

3.1 红黑树的基本概念

TreeMap 的底层数据结构是红黑树（Red - Black Tree），它是一种自平衡的二叉搜索树。红黑树具有以下特性：

• 每个节点要么是红色，要么是黑色。
• 根节点是黑色。
• 每个叶子节点（NIL 节点，空节点）是黑色。
• 如果一个节点是红色的，则它的两个子节点都是黑色的。
• 对每个节点，从该节点到其所有后代叶节点的简单路径上，均包含相同数目的黑色节点。

这些特性保证了红黑树的高度始终保持在 O(log n)，从而使得插入、删除和查找操作的时间复杂度都为 O(log n)。

3.2 红黑树在 TreeMap 中的应用

TreeMap 利用红黑树的特性来存储和管理键值对。每个键值对对应红黑树中的一个节点，键用于比较节点的大小，以确定节点在树中的位置。插入、删除和查找操作都基于红黑树的算法进行，通过比较键的大小来遍历树，找到合适的位置进行操作，并在操作后调整树的结构以保持红黑树的平衡。

3.3 节点结构

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

每个 Entry 节点包含键、值、左右子节点、父节点和颜色信息。

四、核心属性

TreeMap 有几个重要的核心属性，这些属性控制着 TreeMap 的行为和状态：

// 比较器，用于对键进行比较
private final Comparator<? super K> comparator;
// 红黑树的根节点
private transient Entry<K,V> root;
// 键值对的数量
private transient int size = 0;
// 修改次数，用于快速失败机制
private transient int modCount = 0;

• comparator：一个 Comparator 类型的对象，用于对键进行比较。如果为 null，则使用键的自然顺序进行比较。
• root：红黑树的根节点，初始值为 null。
• size：表示 TreeMap 中键值对的数量。
• modCount：记录 TreeMap 的修改次数，用于快速失败机制，当在迭代过程中检测到 modCount 发生变化时，会抛出 ConcurrentModificationException。

五、构造方法

5.1 无参构造方法

public TreeMap() {
    comparator = null;
}

无参构造方法将比较器设置为 null，这意味着 TreeMap 将使用键的自然顺序进行排序。键必须实现 Comparable 接口，否则会抛出 ClassCastException。

5.2 指定比较器的构造方法

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

该构造方法允许用户传入一个自定义的比较器，TreeMap 将使用该比较器对键进行排序。

5.3 从其他 Map 创建 TreeMap 的构造方法

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

此构造方法接受一个 Map 对象作为参数，将该 Map 中的所有键值对添加到新创建的 TreeMap 中，并使用键的自然顺序进行排序。

5.4 从其他 SortedMap 创建 TreeMap 的构造方法

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 对象作为参数，使用该 SortedMap 的比较器对键进行排序，并将其中的键值对添加到新的 TreeMap 中。

import java.util.*;

public class TreeMapConstructors {
    public static void main(String[] args) {
        // 无参构造方法
        TreeMap<String, Integer> treeMap1 = new TreeMap<>();

        // 指定比较器的构造方法
        TreeMap<String, Integer> treeMap2 = new TreeMap<>(Comparator.reverseOrder());

        // 从其他 Map 创建 TreeMap 的构造方法
        Map<String, Integer> hashMap = new HashMap<>();
        hashMap.put("apple", 1);
        hashMap.put("banana", 2);
        TreeMap<String, Integer> treeMap3 = new TreeMap<>(hashMap);

        // 从其他 SortedMap 创建 TreeMap 的构造方法
        SortedMap<String, Integer> sortedMap = new TreeMap<>();
        sortedMap.put("cherry", 3);
        sortedMap.put("date", 4);
        TreeMap<String, Integer> treeMap4 = new TreeMap<>(sortedMap);

        System.out.println("TreeMap1: " + treeMap1);
        System.out.println("TreeMap2: " + treeMap2);
        System.out.println("TreeMap3: " + treeMap3);
        System.out.println("TreeMap4: " + treeMap4);
    }
}

六、常用操作原理

6.1 插入元素（put 方法）

public V put(K key, V value) {
    Entry<K,V> 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<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);
    }
    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. 检查根节点是否为空：如果根节点为空，则创建一个新的节点作为根节点。
2. 查找插入位置：根据比较器或键的自然顺序，从根节点开始遍历红黑树，找到合适的插入位置。
3. 插入新节点：如果找到的位置为空，则创建一个新的节点并插入到该位置。
4. 调整树的平衡：插入新节点后，调用 fixAfterInsertion 方法调整红黑树的结构，以保持其平衡。
5. 更新 size 和 modCount：插入成功后，更新键值对的数量和修改次数。

6.2 获取元素（get 方法）

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

final Entry<K,V> getEntry(Object key) {
    // Offload comparator-based version for sake of performance
    if (comparator != null)
        return getEntryUsingComparator(key);
    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. 根据比较器或键的自然顺序，从根节点开始遍历红黑树。
2. 比较键的大小，根据比较结果决定向左子树还是右子树遍历。
3. 如果找到匹配的键，则返回该节点的值；如果遍历到叶子节点仍未找到，则返回 null。

6.3 删除元素（remove 方法）

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

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

private void deleteEntry(Entry<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) {
        Entry<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.
    Entry<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;
        }
    }
}

remove 方法的主要步骤如下：

1. 查找要删除的节点：调用 getEntry 方法查找要删除的节点。
2. 处理有两个子节点的情况：如果要删除的节点有两个子节点，找到其后继节点，将后继节点的键和值复制到该节点，然后将待删除节点指向后继节点。
3. 处理有一个子节点或没有子节点的情况：将子节点或 null 连接到父节点，并调整树的平衡。
4. 调整树的平衡：删除节点后，调用 fixAfterDeletion 方法调整红黑树的结构，以保持其平衡。
5. 更新 size 和 modCount：删除成功后，更新键值对的数量和修改次数。

6.4 遍历元素

TreeMap 支持多种遍历方式，例如使用 entrySet()、keySet() 和 values() 方法获取相应的集合，然后使用迭代器或增强 for 循环进行遍历。

import java.util.TreeMap;
import java.util.Map;

public class TreeMapTraversal {
    public static void main(String[] args) {
        TreeMap<String, Integer> treeMap = new TreeMap<>();
        treeMap.put("apple", 1);
        treeMap.put("banana", 2);
        treeMap.put("cherry", 3);

        // 遍历键值对
        for (Map.Entry<String, Integer> entry : treeMap.entrySet()) {
            System.out.println(entry.getKey() + ": " + entry.getValue());
        }

        // 遍历键
        for (String key : treeMap.keySet()) {
            System.out.println(key);
        }

        // 遍历值
        for (Integer value : treeMap.values()) {
            System.out.println(value);
        }
    }
}

七、排序机制

7.1 自然排序

如果使用无参构造方法创建 TreeMap，键将按照自然顺序进行排序。键必须实现 Comparable 接口，TreeMap 会调用其 compareTo 方法来比较键的大小。例如，Integer、String 等类都实现了 Comparable 接口。

import java.util.TreeMap;

public class NaturalOrderingExample {
    public static void main(String[] args) {
        TreeMap<Integer, String> treeMap = new TreeMap<>();
        treeMap.put(3, "cherry");
        treeMap.put(1, "apple");
        treeMap.put(2, "banana");
        System.out.println("TreeMap with natural ordering: " + treeMap);
    }
}

7.2 定制排序

如果需要按照自定义的规则对键进行排序，可以使用指定比较器的构造方法。比较器是一个实现了 Comparator 接口的类，需要重写 compare 方法。

import java.util.Comparator;
import java.util.TreeMap;

class Person {
    private String name;
    private int age;

    public Person(String name, int age) {
        this.name = name;
        this.age = age;
    }

    public int getAge() {
        return age;
    }

    @Override
    public String toString() {
        return "Person{name='" + name + "', age=" + age + "}";
    }
}

public class CustomOrderingExample {
    public static void main(String[] args) {
        TreeMap<Person, String> treeMap = new TreeMap<>(Comparator.comparingInt(Person::getAge));
        treeMap.put(new Person("Alice", 25), "Engineer");
        treeMap.put(new Person("Bob", 20), "Student");
        treeMap.put(new Person("Charlie", 30), "Manager");
        System.out.println("TreeMap with custom ordering: " + treeMap);
    }
}

八、性能分析

8.1 时间复杂度

• 插入操作：插入操作的时间复杂度为 O(log n)，因为需要在红黑树中找到合适的位置插入节点，并调整树的结构以保持平衡。
• 删除操作：删除操作的时间复杂度为 O(log n)，同样需要在红黑树中找到要删除的节点，并调整树的结构。
• 查找操作：查找操作的时间复杂度为 O(log n)，可以通过红黑树的特性快速定位节点。

8.2 空间复杂度

TreeMap 的空间复杂度为 O(n)，主要用于存储红黑树的节点。

九、注意事项

9.1 键的可比较性

如果使用自然排序，存储在 TreeMap 中的键必须实现 Comparable 接口，否则会抛出 ClassCastException。如果使用定制排序，则需要提供一个合适的比较器。

9.2 线程安全问题

TreeMap 不是线程安全的。如果在多线程环境下需要使用线程安全的 Map，可以考虑使用 ConcurrentSkipListMap。

9.3 性能考虑

由于 TreeMap 基于红黑树实现，插入、删除和查找操作的时间复杂度为 <math xmlns="http://www.w3.org/1998/Math/MathML"> O ( l o g n ) O(log n) </math>O(logn)，在数据量较大时性能较好。但如果对插入、删除操作的性能要求极高，且不需要排序功能，可以考虑使用 HashMap。

十、总结

TreeMap 是 Java 集合框架中一个非常实用的 Map 实现，它基于红黑树实现了键的有序存储和高效的查找、插入、删除操作。通过自然排序或定制排序，我们可以灵活地对键进行排序。在使用 TreeMap 时，需要注意键的可比较性、线程安全问题和性能考虑。深入理解 TreeMap 的原理和使用方法，有助于我们在实际开发中更好地利用它来处理有序键值对数据。