TreeMap源码解析

一、核心数据结构：红黑树

TreeMap 基于 红黑树（Red-Black Tree） 实现，这是一种自平衡的二叉搜索树，每个节点包含键、值、颜色标记及指向父、左、右子节点的引用：

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

    public K getKey() { return key; }
    public V getValue() { return value; }
    public V setValue(V value) { /* ... */ }
    public boolean equals(Object o) { /* ... */ }
    public int hashCode() { /* ... */ }
    public String toString() { /* ... */ }
}

红黑树的关键性质：

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

二、初始化与比较器

TreeMap 支持两种排序方式：

1. 自然排序：键类必须实现 Comparable 接口。
2. 定制排序：通过构造函数传入 Comparator。

// 自然排序构造函数
public TreeMap() {
    comparator = null; // 使用键的自然顺序
}

// 定制排序构造函数
public TreeMap(Comparator<? super K> comparator) {
    this.comparator = comparator;
}

// 从其他 Map 初始化
public TreeMap(Map<? extends K, ? extends V> m) {
    comparator = null;
    putAll(m); // 内部调用红黑树插入
}

比较器的使用逻辑：

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

三、核心操作源码解析

1. 插入操作：put(K key, V value)

public V put(K key, V value) {
    Entry<K,V> t = root;
    if (t == null) {
        // 根节点直接插入
        compare(key, 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 {
        // 使用自然顺序
        do {
            parent = t;
            cmp = ((Comparable<? super K>)key).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;
}

红黑树调整（插入后修复）：

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 {
            // 父节点是祖父节点的右子节点（对称逻辑）
            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; // 确保根节点为黑色
}

左旋与右旋操作：

// 左旋操作（以 x 为支点）
private void rotateLeft(Entry<K,V> x) {
    if (x != null) {
        Entry<K,V> y = x.right;
        x.right = y.left;
        if (y.left != null)
            y.left.parent = x;
        y.parent = x.parent;
        if (x.parent == null)
            root = y;
        else if (x.parent.left == x)
            x.parent.left = y;
        else
            x.parent.right = y;
        y.left = x;
        x.parent = y;
    }
}

// 右旋操作（对称逻辑）
private void rotateRight(Entry<K,V> y) {
    // 类似左旋，方向相反
}

2. 查询操作：get(Object key)

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

final Entry<K,V> getEntry(Object key) {
    // 根据比较器类型选择查找方式
    if (comparator != null)
        return getEntryUsingComparator(key);
    if (key == null)
        throw new NullPointerException(); // 不支持 null 键

    // 自然顺序查找
    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; // 未找到
}

3. 删除操作：remove(Object key)

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 (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 的引用
        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;
        }
    }
}

// 查找后继节点（中序遍历的下一个节点）
static <K,V> TreeMap.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;
    }
}

四、范围视图与子 Map

TreeMap 提供基于范围的子视图，通过 subMap()、headMap()、tailMap() 方法实现：

// 返回 [fromKey, toKey) 范围的子 Map
public SortedMap<K,V> subMap(K fromKey, K toKey) {
    return new AscendingSubMap<>(
        this, false, fromKey, true, toKey, false);
}

// 返回小于 toKey 的子 Map
public SortedMap<K,V> headMap(K toKey) {
    return new AscendingSubMap<>(
        this, false, null, false, toKey, false);
}

// 返回大于等于 fromKey 的子 Map
public SortedMap<K,V> tailMap(K fromKey) {
    return new AscendingSubMap<>(
        this, false, fromKey, true, null, false);
}

子 Map 实现原理：

• AscendingSubMap 继承自 AbstractMap，内部维护原 TreeMap 的引用和范围边界。
• 所有操作（如 get、put）会检查键是否在范围内，确保子 Map 的修改反映到原 Map。

五、迭代器实现

TreeMap 的迭代器基于红黑树的中序遍历，保证元素按键的升序排列：

// 获取键的迭代器
public Iterator<K> keySet().iterator() {
    return new KeyIterator(getFirstEntry());
}

// 获取第一个节点（最左节点）
final Entry<K,V> getFirstEntry() {
    Entry<K,V> p = root;
    if (p != null)
        while (p.left != null)
            p = p.left;
    return p;
}

// 键迭代器实现
final class KeyIterator extends PrivateEntryIterator<K> {
    KeyIterator(Entry<K,V> first) {
        super(first);
    }
    public K next() {
        return nextEntry().key;
    }
}

// 私有迭代器基类
abstract class PrivateEntryIterator<T> implements Iterator<T> {
    Entry<K,V> next;
    Entry<K,V> lastReturned;
    int expectedModCount;

    PrivateEntryIterator(Entry<K,V> first) {
        expectedModCount = modCount;
        lastReturned = null;
        next = first;
    }

    public final boolean hasNext() {
        return next != null;
    }

    final Entry<K,V> nextEntry() {
        Entry<K,V> e = next;
        if (e == null)
            throw new NoSuchElementException();
        if (modCount != expectedModCount)
            throw new ConcurrentModificationException();

        // 找到下一个节点（中序遍历的后继）
        next = successor(e);
        lastReturned = e;
        return e;
    }

    // 其他方法...
}

六、常见问题解答

1. 为什么 TreeMap 不支持 null 键？

   • 因为 TreeMap 需要通过比较器（或自然顺序）确定键的位置，而 null 无法参与比较（NullPointerException）。

2. TreeMap 的插入 / 删除性能为什么是 O (log n)？

   • 红黑树通过旋转和变色保持平衡，确保树的高度始终为 O (log n)，因此操作复杂度与树高成正比。

3. 如何实现降序遍历？

   NavigableMap<K,V> descendingMap() {
       return new DescendingSubMap<>(this);
   }

4. TreeMap 的空间复杂度是多少？

   • 每个节点需额外存储父节点、左右子节点和颜色标记，空间开销约为 O (n)，略高于 HashMap。

5. TreeMap 与 ConcurrentSkipListMap 的区别？

   • TreeMap：非线程安全，基于红黑树，适用于单线程有序操作。
   • ConcurrentSkipListMap：线程安全，基于跳表，支持高并发有序操作。

总结

TreeMap 通过红黑树实现了 有序映射，其核心优势在于高效的范围查询和有序遍历。源码中最关键的部分是插入 / 删除后的红黑树调整（旋转和变色），这保证了树的平衡性和操作效率。使用时需注意键必须实现 Comparable 接口或提供显式比较器，且不支持 null 键。在多线程场景下，可考虑使用 ConcurrentSkipListMap 替代。

⬅️ TreeMap 🏠 00-Java ➡️ 红黑树