大话红黑树之(2)源码讲解TreeMap-Java

Java TreeMap 源码解析（结合红黑树）

1. 概述

TreeMap 是 Java 中基于红黑树 实现的有序 Map，其核心特性是能保证键的自然顺序，或者根据自定义比较器提供的顺序。通过红黑树的数据结构，TreeMap 保证了常规操作（如插入、删除和查找）的时间复杂度为 O(log n) 。接下来，我们将从源码的角度详细分析 TreeMap 的内部实现，并结合红黑树的特性进行解读。

---

2. TreeMap 中的节点定义

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

    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) {
        V oldValue = this.value;
        this.value = value;
        return oldValue;
    }
}

关键属性说明：

• key：节点的键。

• value：节点的值。

• left 和 right：指向左右子节点。

• parent：指向父节点，方便在树结构中追溯。

• color：节点的颜色属性，红黑树 的核心特性，节点要么是红色要么是黑色。

---

3. TreeMap 插入操作与红黑树调整

TreeMap 中的插入操作通过 put() 方法实现，插入元素后需要通过红黑树的特性来维持平衡。主要的插入逻辑如下：

public V put(K key, V value) {
    Entry<K,V> t = root;
    if (t == null) {
        // 树为空时，创建根节点，颜色设为黑色
        compare(key, key); // type check
        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;

    fixAfterInsertion(e);
    size++;
    modCount++;
    return null;
}

插入操作说明：

• 通过比较器或键的自然顺序找到插入位置。

• 如果插入的新节点已经存在，则更新其值。

• 将新节点插入到合适的位置后，调用 fixAfterInsertion() 方法调整红黑树的平衡性。

---

4. 红黑树的平衡调整

红黑树在插入节点后可能会违反红黑树的性质，TreeMap 中通过 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) {
                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;
}

平衡调整说明：

• 插入节点后，先将其设为红色。

• 如果新节点的父节点也是红色（违反红黑树性质），需要通过旋转（左旋、右旋）和颜色调整修复树的平衡。

• 最终确保根节点始终是黑色。

---

5. 红黑树的旋转操作

旋转是红黑树调整的关键步骤，TreeMap 中定义了 rotateLeft() 和 rotateRight() 两个方法来实现左旋和右旋。

左旋操作：

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

右旋操作：

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

旋转操作说明：

• 左旋 ：节点 p 的右子节点 r 替代 p，p 成为 r 的左子节点。

• 右旋 ：节点 p 的左子节点 l 替代 p，p 成为 l 的右子节点。
  通过这些旋转操作，红黑树能够保持相对平衡，从而维持 O(log n) 的操作效率。

---

6. 总结

TreeMap 通过红黑树来保证数据的有序性和高效操作，其插入、删除等操作都通过红黑树的调整机制（颜色变换与旋转）来保持平衡。通过源码分析，可以看到 TreeMap 内部的红黑树调整过程非常关键，从而确保了查找、插入等操作的高效性。