Java集合-TreeMap

前言

TreeMap 是 Java 集合框架中的一个重要成员，它提供了一种基于键值对的有序存储机制。TreeMap 的核心实现基于红黑树（Red-Black Tree），这是一种自平衡的二叉搜索树，能够在对数时间复杂度 $O(log n)$ 内完成查找、插入和删除等操作。以下将从 TreeMap 的特点、使用场景以及源码实现原理进行详细分析。

特点

TreeMap 具有以下显著特点：

有序性
TreeMap 中的键是按照 自然顺序（Natural Ordering） 或通过自定义的 Comparator 接口指定的顺序进行排序的。这使得键在 TreeMap 中始终保持有序状态。
导航方法
TreeMap 提供了一系列便捷的导航方法，例如 firstKey()、lastKey()、lowerKey()、higherKey() 等，可以根据键的顺序快速定位到特定的键。
子视图支持
TreeMap 支持创建子 Map 视图，例如 subMap()、headMap()、tailMap()，这些视图允许操作特定范围内的键值对。
线程安全性
TreeMap 不是线程安全的。如果在多线程环境下使用，需要通过外部同步机制（如 Collections.synchronizedMap）来保证线程安全。
键值限制
TreeMap 不允许键为 null，因为 null 无法与其他对象进行比较。
TreeMap 允许值为 null。

使用场景

由于 TreeMap 的键具有有序性，它特别适用于以下场景：

实现需要按键排序的字典或电话簿。
需要快速查找并按顺序输出的应用程序。
对键值对进行范围操作的场景，例如获取某个范围内的数据。

源码分析

数据结构

TreeMap 的底层数据结构是红黑树，一种自平衡的二叉搜索树。红黑树通过特定的性质保证树的平衡，从而确保所有主要操作的时间复杂度为 $O(log n)$。

红黑树节点定义

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; // 节点颜色（红或黑） true 表示黑色，false 表示红色
    // 构造函数和其他方法...
}

每个 Entry 节点包含键、值、左右子节点、父节点以及颜色属性（BLACK 或 RED）。

红黑树的性质

红黑树通过以下性质维持平衡：

节点颜色：每个节点要么是红色，要么是黑色。
根节点：根节点始终是黑色。
叶子节点：所有叶子节点（NIL 节点）是黑色。
红色节点限制：红色节点的两个子节点必须是黑色（即红色节点不能相邻）。
黑色高度：从任一节点到其每个叶子节点的所有路径包含相同数量的黑色节点。

这些性质保证了红黑树的高度不会超过 $2log(n+1)$，从而保持操作的高效性。

关键属性

1
2
3

private final Comparator<? super K> comparator; // 自定义比较器
private transient Entry<K,V> root;             // 红黑树根节点
private transient int size = 0;                // 元素数量

核心方法

以下从源码角度分析 TreeMap 的核心操作：插入、删除、查找和导航方法。

插入操作（put）

put 方法用于向 TreeMap 中插入键值对。插入过程包括查找位置、插入节点和调整红黑树。

插入步骤

查找插入位置：从根节点开始，根据键的比较结果（通过 compareTo 或自定义 Comparator）向左或向右遍历，直到找到 null 位置。
插入新节点：新节点的初始颜色为红色。
调整红黑树：插入红色节点可能破坏红黑树的性质，需要通过旋转和颜色调整来修复。

调整场景

父节点为黑色：直接插入，无需调整。
父节点为红色：
- 叔叔节点为红色：将父节点和叔叔节点变为黑色，祖父节点变为红色，然后以祖父节点为当前节点继续调整。
- 叔叔节点为黑色或 null：通过左旋或右旋调整树结构，并更新颜色。

源码

插入操作的源码如下：

public V put(K key, V value) {
    Entry<K,V> t = root;
    if (t == null) {
        // 根节点为空时，直接创建新节点
        root = new Entry<>(key, value, null);
        size = 1;
        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 {
        // 使用自然顺序查找插入位置
        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);
    }
    // 创建新节点并插入到红黑树中
    Entry<K,V> e = new Entry<>(key, value, parent);
    if (cmp < 0)
        parent.left = e;
    else
        parent.right = e;
    fixAfterInsertion(e); // 调整红黑树平衡
    size++;
    return null;
}

红黑树平衡调整 fixAfterInsertion(Entry<K,V> x)：

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) {
                // 叔叔节点是红色：颜色翻转
                setColor(parentOf(x), BLACK);
                setColor(y, BLACK);
                setColor(parentOf(parentOf(x)), RED);
                x = parentOf(parentOf(x));
            } else {
                if (x == rightOf(parentOf(x))) {
                    // 左旋父节点
                    x = parentOf(x);
                    rotateLeft(x);
                }
                // 右旋祖父节点并调整颜色
                setColor(parentOf(x), BLACK);
                setColor(parentOf(parentOf(x)), RED);
                rotateRight(parentOf(parentOf(x)));
            }
        } else {
            // 对称逻辑（右子树处理）
        }
    }
    root.color = BLACK; // 根节点始终为黑色
}

删除操作（remove）

remove 方法用于删除指定键的节点。删除后需要调整红黑树以保持其性质。

删除步骤

查找节点：如果目标节点有两个子节点，则找到其后继节点（右子树的最小节点），交换键值对，然后删除后继节点。
删除节点：根据节点的子节点情况执行删除。
调整红黑树：如果删除的是黑色节点，可能破坏性质，需要调整。

调整场景

兄弟节点为红色：通过旋转和颜色调整修复。
兄弟节点为黑色：根据兄弟节点的子节点颜色进行不同调整。

源码

删除操作的源码如下：

private void deleteEntry(Entry<K,V> p) {
    // 找到替代节点（若左右子节点均存在）
    if (p.left != null && p.right != null) {
        Entry<K,V> s = successor(p); // 找到后继节点
        p.key = s.key;
        p.value = s.value;
        p = s;
    }
    // 执行删除并修复红黑树
    Entry<K,V> replacement = (p.left != null ? p.left : p.right);
    if (replacement != null) {
        // 处理子节点替换逻辑
        // ...
        fixAfterDeletion(replacement); // 修复删除后的平衡
    } else if (p.parent == null) {
        root = null;
    } else {
        fixAfterDeletion(p); // 修复删除后的平衡
        if (p.parent != null) {
            if (p == p.parent.left)
                p.parent.left = null;
            else
                p.parent.right = null;
        }
    }
    size--;
}

查找操作（get）

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) {
    Entry<K,V> p = root;
    while (p != null) {
        int cmp = compare(key, p.key);
        if (cmp < 0)
            p = p.left;
        else if (cmp > 0)
            p = p.right;
        else
            return p;
    }
    return null;
}

导航方法

TreeMap 提供了丰富的导航方法，利用红黑树的有序性快速定位键。例如：

firstKey()：返回最小键。
lastKey()：返回最大键。

源码如下：

public K firstKey() {
    return key(getFirstEntry());
}

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

public K lastKey() {
    return key(getLastEntry());
}

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

后记

TreeMap 是 Java 集合框架中一个基于红黑树实现的有序 Map，提供 $O(log n) $的查找、插入和删除操作。它支持导航方法和子视图，适用于需要保持键有序性的场景。