TreeMap / 红黑树源码

概述

TreeMap 是基于红黑树（Red-Black Tree）实现的有序映射表。它按照键的自然顺序（Comparable）或构造器传入的 Comparator 排序，保证所有操作的时间复杂度为 O(log n)。

TreeMap 实现了 NavigableMap 接口，提供了丰富的有序操作：范围查找、逆序视图、最近匹配等。

本文基于 OpenJDK 21 源码，从红黑树结构开始，逐层深入到搜索树插入、红黑树平衡修正、节点删除旋转以及有序操作等核心实现。

本文基于 OpenJDK 21 源码分析。

---

核心源码解析

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

    private static final boolean RED   = false;
    private static final boolean BLACK = true;
}

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; // 默认黑色
}

红黑树 5 条性质：

1. 每个节点要么是红色，要么是黑色
2. 根节点是黑色
3. 叶子节点（null）是黑色
4. 红色节点的子节点必须是黑色（不能有两个连续的红色节点）
5. 任意节点到其每个叶子节点的路径上都包含相同数量的黑色节点

② put(K, V) 的二叉搜索树插入

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 自然排序
        Comparable<? super K> k = (Comparable<? super K>) key;
        do {
            parent = t;
            cmp = k.compareTo(t.key);
            // ... 同上
        } 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;
}

• 从根节点开始，通过 compare() 比较找到插入位置
• 键已存在则替换值
• 插入后调用 fixAfterInsertion() 修正红黑树平衡

③ fixAfterInsertion(Entry) 的红黑树修正

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);      // 情况1：叔父为红
                setColor(y, BLACK);
                setColor(parentOf(parentOf(x)), RED);
                x = parentOf(parentOf(x));
            } else {
                if (x == rightOf(parentOf(x))) {
                    x = parentOf(x);
                    rotateLeft(x);                  // 情况2：左旋
                }
                setColor(parentOf(x), BLACK);      // 情况3：右旋
                setColor(parentOf(parentOf(x)), RED);
                rotateRight(parentOf(parentOf(x)));
            }
        } else {
            // 对称操作（父节点在右侧）
            // ...
        }
    }
    root.color = BLACK; // 根节点保持黑色
}

三种情况：

1. 叔父为红色：父和叔变黑，祖父变红，x 上移
2. 叔父为黑色，x 在内侧：对父节点左旋，转为情况 3
3. 叔父为黑色，x 在外侧：父变黑，祖父变红，对祖父右旋

④ deleteEntry(Entry) 的节点删除

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

    // 有两个子节点 → 找后继节点替换
    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) {
        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.left = p.right = p.parent = null; // 清理引用
        if (p.color == BLACK)
            fixAfterDeletion(replacement);
    } else if (p.parent == null) {
        root = null;
    } else {
        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(p) 找后继节点：

static <K,V> 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;
    }
}

• 有右子节点 → 右子树最左节点
• 无右子节点 → 向上找第一个作为左子节点的父节点

⑤ rotateLeft(Entry) 左旋

private void rotateLeft(Entry<K,V> p) {
    if (p != null) {
        Entry<K,V> r = p.right;
        p.right = r.left;
        if (r.left != null)
            r.left.parent = p;
        r.parent = p.parent;
        if (p.parent == null)
            root = r;
        else if (p.parent.left == p)
            p.parent.left = r;
        else
            p.parent.right = r;
        r.left = p;
        p.parent = r;
    }
}

⑥ rotateRight(Entry) 右旋

private void rotateRight(Entry<K,V> p) {
    if (p != null) {
        Entry<K,V> l = p.left;
        p.left = l.right;
        if (l.right != null)
            l.right.parent = p;
        l.parent = p.parent;
        if (p.parent == null)
            root = l;
        else if (p.parent.right == p)
            p.parent.right = l;
        else
            p.parent.left = l;
        l.right = p;
        p.parent = l;
    }
}

⑦ Comparator vs Comparable 的排序

public TreeMap() {
    comparator = null; // 使用键的 Comparable 自然排序
}

public TreeMap(Comparator<? super K> comparator) {
    this.comparator = comparator; // 自定义比较器
}

final int compare(Object k1, Object k2) {
    return comparator == null ? ((Comparable<? super K>) k1).compareTo((K) k2)
                              : comparator.compare((K) k1, (K) k2);
}

• compare() 统一入口，首次调用时还用于类型检查
• 两种排序方式的选择在 put() 中通过 comparator != null 分支实现

⑧ NavigableMap 的范围查找

public NavigableMap<K,V> subMap(K fromKey, boolean fromInclusive,
                                 K toKey,   boolean toInclusive) {
    return new AscendingSubMap<>(this,
                                 false, fromKey, fromInclusive,
                                 false, toKey,   toInclusive);
}

public NavigableMap<K,V> headMap(K toKey, boolean inclusive) {
    return new AscendingSubMap<>(this,
                                 true,  null,  true,
                                 false, toKey, inclusive);
}

public NavigableMap<K,V> tailMap(K fromKey, boolean inclusive) {
    return new AscendingSubMap<>(this,
                                 false, fromKey, inclusive,
                                 true,  null,    true);
}

NavigableSubMap 是一个抽象内部类，AscendingSubMap 和 DescendingSubMap 分别是升序和降序视图：

• 范围视图上的操作直接委托给原 TreeMap
• 超出范围的 put() 会抛 IllegalArgumentException
• 范围视图支持所有 NavigableMap 操作

---

总结

TreeMap 的红黑树实现提供了 O(log n) 的有序映射操作。其核心设计包括：

• 红黑树自平衡：通过颜色标记和左/右旋转维持 O(log n) 的高度
• 自然顺序/比较器：支持 Comparable 和 Comparator 两种排序方式
• 有序操作：NavigableMap 提供范围查找、最近匹配、逆序视图等能力
• O(log n) 保证：所有核心操作（put、get、remove）时间复杂度为 O(log n)

与 HashMap 相比，TreeMap 牺牲了常数时间的随机访问性能，但提供了有序集合功能，适用于需要排序或范围查询的场景。