TreeMap 源码分析

最新推荐文章于 2025-06-21 11:54:54 发布

原创最新推荐文章于 2025-06-21 11:54:54 发布 · 790 阅读

11 ·

CC 4.0 BY-SA版权

文章标签：

#算法 #数据结构

JavaSE 同时被 2 个专栏收录

72 篇文章

订阅专栏

Java

72 篇文章

订阅专栏

TreeMap 源码分析

TreeMap 是 Java 集合框架中基于红黑树实现的有序键值对集合，它不仅能够按键的自然顺序或用户自定义的比较器顺序存储元素，还能保证插入、删除和查找操作的时间复杂度为 O(log n)。本文将从源码角度深入分析 TreeMap 的实现机制。

1. TreeMap 简介

TreeMap 是 Java 集合框架中的一部分，它基于红黑树（Red-Black Tree）实现，用来存储键值对并按键进行排序。TreeMap 与 HashMap 不同，它保证了键的顺序性，因此可以高效地执行有序集合的操作，例如获取最小值、最大值和范围查询。

红黑树的主要特点是：

它是一种 自平衡的二叉搜索树，保证最坏情况下的查找、插入、删除操作的时间复杂度为 O(log n)。
通过节点的颜色（红色或黑色）来维持平衡，红黑树的高度近似于 log(n)，其中 n 是树中的节点数量。

2. TreeMap 的基本构造函数

首先看一下 TreeMap 的构造函数，它允许我们选择使用自然顺序或者指定的 Comparator 进行排序。

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

    // 用于比较键的比较器，null 表示自然顺序
    private final Comparator<? super K> comparator;

    // 构造一个空的 TreeMap，使用键的自然顺序
    public TreeMap() {
        comparator = null;
    }

    // 构造一个空的 TreeMap，使用指定的比较器
    public TreeMap(Comparator<? super K> comparator) {
        this.comparator = comparator;
    }
    
    // 其他构造方法省略...
}

这里的 comparator 是一个用来比较键的比较器。如果 comparator 为 null，则使用键的自然顺序（如 Integer、String 等类的 compareTo() 方法）。

3. TreeMap 的核心：红黑树节点

TreeMap 的实现依赖于内部类 Entry，该类用于表示树中的每个节点。Entry 继承了 Map.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; // 节点颜色，默认黑色
    
    // 构造方法
    Entry(K key, V value, Entry<K,V> parent) {
        this.key = key;
        this.value = value;
        this.parent = parent;
    }

    // 实现 Map.Entry 的接口方法
    public K getKey() { return key; }
    public V getValue() { return value; }
    public V setValue(V value) {
        V oldValue = this.value;
        this.value = value;
        return oldValue;
    }
}

在这个 Entry 类中，除了 key 和 value，还包含指向左、右子节点和父节点的指针，以及用于表示节点颜色的 color 字段。红黑树的平衡性通过维护这些节点及其颜色来实现。

4. 插入操作

在 TreeMap 中，键值对的插入操作是通过 put() 方法实现的。插入过程类似于二叉搜索树的插入，插入后可能会打破红黑树的性质，因此需要调用 fixAfterInsertion() 来进行调整。

public V put(K key, V value) {
    Entry<K,V> t = root;
    
    // 如果根节点为空，插入的节点为根节点
    if (t == null) {
        compare(key, key); // 检查键的可比性
        root = new Entry<>(key, value, null);
        size = 1;
        modCount++;
        return null;
    }

    int cmp;
    Entry<K,V> parent;
    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 { // 使用键的自然顺序
        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;

    size++;
    modCount++;
    
    // 修复红黑树平衡性
    fixAfterInsertion(e);
    return null;
}

插入逻辑：

首先检查 root 是否为 null，如果 root 为 null，说明这是一个空树，将新节点设为根节点。

否则，从根节点开始，通过 compareTo 或者 Comparator 比较键的大小，找到要插入的位置。

找到合适的插入位置后，创建一个新的 Entry 节点并插入树中。

插入节点后，调用 fixAfterInsertion() 方法来调整红黑树的平衡性。

5. 红黑树插入后的调整：fixAfterInsertion

红黑树的性质包括：

节点是红色或黑色。
根节点是黑色。
叶节点（null 节点）是黑色。
红色节点的子节点必须是黑色（不能有两个连续的红色节点）。
从任一节点到其每个叶节点的所有路径都必须包含相同数量的黑色节点。

fixAfterInsertion() 方法用于调整插入后打破红黑树性质的情况。主要包括三种情况：

父节点是黑色，树仍然平衡，无需调整。
父节点是红色，可能需要通过旋转和着色来恢复红黑树的平衡。

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) {
                // 情况 1: 叔叔节点为红色
                setColor(parentOf(x), BLACK);
                setColor(y, BLACK);
                setColor(parentOf(parentOf(x)), RED);
                x = parentOf(parentOf(x));
            } else {
                if (x == rightOf(parentOf(x))) {
                    // 情况 2: 叔叔为黑色，且当前节点是右孩子
                    x = parentOf(x);
                    rotateLeft(x);
                }
                // 情况 3: 叔叔为黑色，且当前节点是左孩子
                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;
}

6. TreeMap 的删除操作

TreeMap 中删除节点的操作较为复杂，因为删除可能破坏红黑树的平衡性，尤其是删除红黑树中的黑色节点时。

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

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