TreeMap 源码解析-CFANZ编程社区

1、简述

今天我们来介绍下 TreeMap，TreeMap 是基于红黑树结构实现的一种 Map。
那什么又是红黑数呢？
（1）每个节点必须是红色或者黑色；
（2）根节点是黑色；
（3）每个叶子节点（NIL 节点，空节点）是黑色的；
（4）如果一个节点是红色的，则它两个字节点都是黑色的，也就是说在一条路径上不能出现相邻的两个红色节点；
（5）从任一节点到每个叶子的所有路径都包含相同数目的黑色节点。
如何保持红黑树特性？
每当添加或删除节点后，红黑树就发生了变化，可能不再满足上述五种特性，所以为了保持红黑树以上特性，就需要三个动作：左旋、右旋、着色。

2、归纳

继承了 AbstractMap，实现了 NavigableMap 接口，所以支持一系列的导航方法，实现了 Cloneable 和 Serializable 接口，所以支持复制（拷贝）、序列化。
可在构造函数传入自己的比较器。
其底层是基于红黑树实现的，无序，不可重复，不允许 null 键，但允许 null 值存在。
单线程安全，多线程不安全。

3、分析

3.1、接口

在分析 TreMap 源码之前，我们先来看看集合的接口 Map。

public interface Map<K, V> {
    ...
    // 增
    V put(K key, V value);
    // 删
    V remove(Object key);
    // 查
    V get(Object key);
    ...
}

在上述接口中，我只抽取了比较重要的几个方法，然后以此为后续重点分析目标，其 Map 接口对应的源码中远不止上述几个方法，有兴趣的同学可以自行翻阅。

3.2、静态内部类

public class TreeMap<K, V> extends AbstractMap<K, V>
        implements NavigableMap<K, V>, Cloneable, java.io.Serializable {
    ...
    static final class TreeMapEntry<K, V> implements Map.Entry<K, V> {
        K key;
        V value;
        // 左节点
        TreeMapEntry<K, V> left;
        // 右节点
        TreeMapEntry<K, V> right;
        // 父节点
        TreeMapEntry<K, V> parent;
        // 每个节点的颜色，红黑树特性，红色或者黑色
        boolean color = BLACK;

        /**
         * 有参构造函数
         *
         * @param key    键
         * @param value  元素值
         * @param parent 父节点
         */
        TreeMapEntry(K key, V value, TreeMapEntry<K, V> parent) {
            this.key = key;
            this.value = value;
            this.parent = parent;
        }
        ...
    }
    ...
}

3.3、成员变量

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

    // 这是一个比较器，方便插入查找元素等操作
    private final Comparator<? super K> comparator;

    // 红黑树的根节点：每个节点是一个 TreeMapEntry
    private transient TreeMapEntry<K, V> root;

    // 红黑树中元素的个数
    private transient int size = 0;

    // 对红黑树结构修改的次数
    private transient int modCount = 0;
    ...
}

3.4、构造函数

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

    /**
     * 无参构造函数
     * 默认，且比较器为空
     */
    public TreeMap() {
        comparator = null;
    }

    /**
     * 有参构造函数
     *
     * @param comparator 指定一个比较器
     */
    public TreeMap(Comparator<? super K> comparator) {
        this.comparator = comparator;
    }

    /**
     * 有参构造函数
     *
     * @param m 指定一个 map 创建，比较器为空，元素自然排序
     */
    public TreeMap(Map<? extends K, ? extends V> m) {
        comparator = null;
        putAll(m);
    }

    /**
     * 有参构造函数
     *
     * @param m 指定 SortedMap，根据 SortedMap 的比较器来来维持 TreeMap 的顺序
     */
    public TreeMap(SortedMap<K, ? extends V> m) {
        comparator = m.comparator();
        try {
            buildFromSorted(m.size(), m.entrySet().iterator(), null, null);
        } catch (java.io.IOException cannotHappen) {
        } catch (ClassNotFoundException cannotHappen) {
        }
    }
    ...
}

3.5、增操作

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

    /**
     * @param key
     * @param value
     * @return
     */
    public V put(K key, V value) {
        TreeMapEntry<K, V> t = root;
        // 如果 root 为 null 说明是添加第一个元素，直接实例化一个 TreeMapEntry 赋值给 root
        if (t == null) {
            // 类型（可能为空）检查
            compare(key, key);
            root = new TreeMapEntry<>(key, value, null);
            size = 1;
            modCount++;
            return null;
        }
        int cmp;
        // 如果 root 不为 null，说明已存在元素
        TreeMapEntry<K, V> parent;
        // 拆分比较器和可比较路径
        Comparator<? super K> cpr = comparator;
        if (cpr != null) {// 如果比较器不为 null 则使用比较器
            do {// 找到元素的插入位置
                parent = t;// parent 赋值
                cmp = cpr.compare(key, t.key);
                if (cmp < 0)// 当前 key 小于节点 key 向左子树查找
                    t = t.left;
                else if (cmp > 0)// 当前 key 大于节点 key 向右子树查找
                    t = t.right;
                else// 相等的情况下直接更新节点值
                    return t.setValue(value);
            } while (t != null);
        } else {// 如果比较器为 null 则使用默认比较器
            if (key == null)// 如果 key 为 null 则抛出异常
                throw new NullPointerException();
            @SuppressWarnings("unchecked")
            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);
        }
        // 定义一个新的节点
        TreeMapEntry<K, V> e = new TreeMapEntry<>(key, value, parent);
        // 根据比较结果决定插入到左子树还是右子树
        if (cmp < 0)
            parent.left = e;
        else
            parent.right = e;
        fixAfterInsertion(e);// 保持红黑树性质，插入后进行修正
        size++;// 元素树自增
        modCount++;// 元素树结构发生变化自增
        return null;
    }

    /**
     * 插入元素后，保持红黑树性质（插入后进行红黑树修正）
     *
     * @param x 插入的节点
     */
    private void fixAfterInsertion(TreeMapEntry<K, V> x) {
        // 将新插入节点的颜色设置为红色
        x.color = RED;
        // while 循环，保证新插入节点 x 不是根节点或者新插入节点 x 的父节点不是红色（这两种情况不需要调整）
        while (x != null && x != root && x.parent.color == RED) {
            // 如果新插入节点 x 的父节点是祖父节点的左孩子
            if (parentOf(x) == leftOf(parentOf(parentOf(x)))) {
                // 取得新插入节点 x 的叔叔节点
                TreeMapEntry<K, V> y = rightOf(parentOf(parentOf(x)));
                // 如果新插入 x 的父节点是红色
                if (colorOf(y) == RED) {
                    // 将 x 的父节点设置为黑色
                    setColor(parentOf(x), BLACK);
                    // 将 x 的叔叔节点设置为黑色
                    setColor(y, BLACK);
                    // 将 x 的祖父节点设置为红色
                    setColor(parentOf(parentOf(x)), RED);
                    // 将 x 指向祖父节点，如果 x 的祖父节点的父节点是红色，按照上面的步奏继续循环
                    x = parentOf(parentOf(x));
                } else {
                    // 如果新插入 x 的叔叔节点是黑色或缺少，且 x 的父节点是祖父节点的右孩子
                    if (x == rightOf(parentOf(x))) {
                        // 左旋父节点
                        x = parentOf(x);
                        rotateLeft(x);
                    }
                    // 如果新插入 x 的叔叔节点是黑色或缺少，且 x 的父节点是祖父节点的左孩子
                    // 将 x 的父节点设置为黑色
                    setColor(parentOf(x), BLACK);
                    // 将 x 的祖父节点设置为红色
                    setColor(parentOf(parentOf(x)), RED);
                    // 右旋 x 的祖父节点
                    rotateRight(parentOf(parentOf(x)));
                }
            } else { // 如果新插入节点 x 的父节点是祖父节点的右孩子和上面的相似
                TreeMapEntry<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;
    }
    ...
}

3.6、删操作

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

    /**
     * 根据 key 进行删操作
     *
     * @param key 键
     * @return 返回删除后的元素值
     */
    public V remove(Object key) {
        // 根据 key 查找到对应的节点对象
        TreeMapEntry<K, V> p = getEntry(key);
        if (p == null)
            return null;

        // 记录 key 对应的 value，供返回使用
        V oldValue = p.value;
        // 删除节点
        deleteEntry(p);
        return oldValue;
    }

    /**
     * 删除节点
     *
     * @param p 节点
     */
    private void deleteEntry(TreeMapEntry<K, V> p) {
        modCount++;
        // 元素个数减一
        size--;
        // 如果被删除的节点 p 的左孩子和右孩子都不为空，则查找其替代节
        if (p.left != null && p.right != null) {
            // 查找 p 的替代节点
            TreeMapEntry<K, V> s = successor(p);
            p.key = s.key;
            p.value = s.value;
            p = s;
        }
        TreeMapEntry<K, V> replacement = (p.left != null ? p.left : p.right);
        if (replacement != null) {
            // 将 p 的父节点拷贝给替代节点
            replacement.parent = p.parent;
            // 如果替代节点 p 的父节点为空，也就是 p 为跟节点，则将 replacement 设置为根节点
            if (p.parent == null)
                root = replacement;
                // 如果替代节点 p 是其父节点的左孩子，则将 replacement 设置为其父节点的左孩子
            else if (p == p.parent.left)
                p.parent.left = replacement;
                // 如果替代节点 p 是其父节点的左孩子，则将 replacement 设置为其父节点的右孩子
            else
                p.parent.right = replacement;
            // 将替代节点 p 的 left、right、parent 的指针都指向空
            p.left = p.right = p.parent = null;
            // 如果替代节点 p 的颜色是黑色，则需要调整红黑树以保持其平衡
            if (p.color == BLACK)
                fixAfterDeletion(replacement);
        } else if (p.parent == null) {
            // 如果要替代节点 p 没有父节点，代表 p 为根节点，直接删除即可
            root = null;
        } else {
            // 如果 p 的颜色是黑色，则调整红黑树
            if (p.color == BLACK)
                fixAfterDeletion(p);
            // 下面删除替代节点 p
            if (p.parent != null) {
                // 解除 p 的父节点对 p 的引用
                if (p == p.parent.left)
                    p.parent.left = null;
                else if (p == p.parent.right)
                    p.parent.right = null;
                // 解除 p 对 p 父节点的引用
                p.parent = null;
            }
        }
    }
    ...
}

3.7、查操作

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

    /**
     * 根据 key 进行查操作
     *
     * @param key 键
     * @return 返回查找后的元素值，如若不存在返回 null
     */
    public V get(Object key) {
        TreeMapEntry<K, V> p = getEntry(key);
        return (p == null ? null : p.value);
    }

    /**
     * 根据 key 获取对应的节点
     *
     * @param key 键
     * @return 返回节点
     */
    final TreeMapEntry<K, V> getEntry(Object key) {
        // 如果比较器为空，只是用 key 作为比较器查询
        if (comparator != null)
            return getEntryUsingComparator(key);
        if (key == null)
            throw new NullPointerException();
        Comparable<? super K> k = (Comparable<? super K>) key;
        // 取得 root 节点
        TreeMapEntry<K, V> p = root;
        // 核心来了：从 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;
    }
    ...
}