TreeSet的add方法源码分析

一、JDK 1.8 中 TreeSet 的 add 方法源码详细分析

TreeSet 是 Java 集合框架中的一个有序集合类，基于红黑树（TreeMap）实现。TreeSet 的 add 方法用于向集合中添加元素。如果元素已存在，则不会添加并返回 false；如果元素不存在，则添加并返回 true。下面我们将从 TreeSet 的 add 方法入手，逐步分析其源码实现

二、TreeSet 的 add 方法

TreeSet 的 add 方法源码如下:

参数说明

• e：要添加的元素。

• PRESENT：一个固定的 Object 对象，作为 TreeMap 的值（value）

返回值

• 如果元素 e 不存在，map.put(e, PRESENT) 返回 null，表示添加成功

• 如果元素 e 已存在，map.put(e, PRESENT) 返回 PRESENT，表示添加失败

三、TreeSet 的内部实现

TreeSet 的内部是基于 TreeMap 实现的。TreeSet 的元素作为 TreeMap 的键（key），而值（value）是一个固定的 PRESENT 对象（private static final Object PRESENT = new Object()）

TreeSet 的构造函数

TreeSet 的构造函数会初始化内部的 TreeMap。例如：

PRESENT 对象

PRESENT 是一个静态常量，用于作为 TreeMap 的值：

四、TreeMap 的 put 方法

TreeSet 的 add 方法依赖于 TreeMap 的 put 方法。以下是 TreeMap 的 put 方法的源码（JDK 1.8）：

    public V put(K key, V value) {
        Entry<K,V> t = root; // 获取根节点
        if (t == null) { // 如果树为空
            compare(key, key); // 检查键是否为 null
            root = new Entry<>(key, value, null); // 创建根节点
            size = 1; // 大小设置为 1
            modCount++; // 修改计数器加 1
            return null; // 返回 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) // 如果键为 null
                throw new NullPointerException(); // 抛出异常
            Comparable<? super K> k = (Comparable<? super K>) key; // 强制转换为 Comparable
            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++; // 大小加 1
        modCount++; // 修改计数器加 1
        return null; // 返回 null，表示添加成功
    }

五、源码解析

1、检查树是否为空：

• 如果树为空，创建根节点并返回 null

2、比较键：

• 如果使用自定义比较器，调用 compare 方法比较键

• 如果使用自然顺序，调用 compareTo 方法比较键

3、查找插入位置：

• 根据比较结果，移动到左子树或右子树

• 如果找到相同的键，更新值并返回旧值

4、插入新节点：

• 创建新节点，并根据比较结果将其作为左子节点或右子节点

5、修复红黑树的性质：

• 调用 fixAfterInsertion 方法修复红黑树的性质（如颜色调整和旋转）

6、更新大小和修改计数器：

• 大小加 1

• 修改计数器加 1

六、总结

TreeSet 的 add 方法通过调用内部 TreeMap 的 put 方法实现元素的添加。它的核心逻辑是：

• 如果元素不存在，添加到 TreeMap 中并返回 true

• 如果元素已存在，返回 false

在 JDK 1.8 中，TreeMap 的实现基于红黑树，保证插入、删除和查找操作的时间复杂度为 O(log n)。TreeSet 的 add 方法适合需要有序集合的场景。在使用时，注意实现 Comparable 接口或提供 Comparator，并根据需求选择合适的集合类