TreeSet源码分析

一、类定义与核心属性

1. 类声明

public class TreeSet<E> extends AbstractSet<E>
    implements NavigableSet<E>, Cloneable, java.io.Serializable {
    // 底层依赖 TreeMap 实现
    private transient NavigableMap<E, Object> m;
    // 占位值，作为 TreeMap 的 value
    private static final Object PRESENT = new Object();
    // 构造函数传入的 Comparator（可能为 null，代表自然排序）
    private transient Comparator<? super E> comparator;
}

2. 关键属性解析

• NavigableMap m：

  TreeSet 的所有操作均委托给 TreeMap（红黑树实现的有序映射）。元素作为 TreeMap 的键（Key）存储，值（Value）统一为静态常量 PRESENT（避免额外内存开销）。

• Comparator comparator：

  用于定制排序规则。若为 null，则要求元素实现 Comparable 接口（自然排序）。

• PRESENT：

  占位对象，所有元素在 TreeMap 中对应的值均为该对象，仅用于标记键存在。

二、构造函数源码分析

1. 无参构造（自然排序）

public TreeSet() {
    this.comparator = null;
    // 底层创建 TreeMap（自然排序）
    this.m = new TreeMap<>();
}

2. 传入 Comparator 的构造函数

public TreeSet(Comparator<? super E> comparator) {
    this.comparator = comparator;
    // 底层创建带 Comparator 的 TreeMap
    this.m = new TreeMap<>(comparator);
}

3. 其他构造函数

• TreeSet(Collection c)：

  先将集合元素转为数组，再按自然排序或 Comparator 排序后插入。

• TreeSet(SortedSet s)：

  直接使用 SortedSet 的 comparator 初始化底层 TreeMap。

三、核心方法源码解析

1. add(E e)：元素插入

public boolean add(E e) {
    // 调用 TreeMap 的 put 方法，键为 e，值为 PRESENT
    return m.put(e, PRESENT) == null;
}

• 核心逻辑：

  通过 TreeMap.put(e, PRESENT) 实现插入。由于 TreeMap 的键不允许重复，若 e 已存在，put 返回已存在的值（非 null），因此 add 返回 false（去重逻辑）；若不存在，返回 true。

• 红黑树操作：

  TreeMap.put 内部通过红黑树的插入算法（含旋转和颜色调整）保证树平衡，时间复杂度为 O(log n)。

2. contains(Object o)：元素查询

public boolean contains(Object o) {
    // 调用 TreeMap 的 containsKey 方法
    return m.containsKey(o);
}

• 核心逻辑：

  直接查询 TreeMap 中是否存在键 o，红黑树查找时间复杂度为 O(log n)。

3. remove(Object o)：元素删除

public boolean remove(Object o) {
    // 调用 TreeMap 的 remove 方法
    return m.remove(o) == PRESENT;
}

• 核心逻辑：

  删除 TreeMap 中键为 o 的条目，若存在则返回 true。红黑树删除后会自动平衡，时间复杂度为 O(log n)。

4. iterator()：迭代器

public Iterator<E> iterator() {
    // 返回 TreeMap 键的迭代器（按排序顺序）
    return m.navigableKeySet().iterator();
}

• 迭代顺序：

  红黑树的中序遍历顺序（即排序后的顺序）。

5. 范围查询方法（如 subSet(E fromElement, E toElement)）

public NavigableSet<E> subSet(E fromElement, boolean fromInclusive,
                             E toElement, boolean toInclusive) {
    // 调用 TreeMap 的 subMap 方法，返回子树视图
    return new TreeSet<>(m.subMap(fromElement, fromInclusive,
                                  toElement, toInclusive));
}

• 底层实现：

  通过 TreeMap.subMap 获取子树范围，返回新的 TreeSet（包装子树视图），操作直接影响原集合。

四、红黑树与排序逻辑

1. 元素比较逻辑

TreeMap 在插入 / 查询元素时，通过 compare 方法确定元素位置：

// TreeMap 内部比较方法（简化版）
private int compare(Object k1, Object k2) {
    // 若存在 Comparator，使用定制比较
    if (comparator != null) {
        return comparator.compare((E) k1, (E) k2);
    }
    // 否则要求元素实现 Comparable
    Comparable<? super E> cpr = (Comparable<? super E>) k1;
    return cpr.compareTo((E) k2);
}

• 异常处理：

  若元素未实现 Comparable 且未传入 Comparator，调用 compare 时会抛出 ClassCastException。

2. 红黑树特性

• 自平衡机制：

  通过以下规则保证树高平衡（最长路径不超过最短路径的 2 倍）：

  1. 每个节点是红色或黑色。
  2. 根节点是黑色。
  3. 叶子节点（NIL 节点）是黑色。
  4. 红色节点的子节点必须是黑色（避免连续红色节点）。
  5. 从任一节点到其每个叶子节点的所有路径包含相同数量的黑色节点。

• 旋转操作：

  插入 / 删除后通过左旋（rotateLeft）和右旋（rotateRight）调整树结构，配合颜色翻转保持平衡。

五、与 TreeMap 的关联

1. 存储关系

2. 视图机制

范围查询方法（如 subSet、headSet）返回的 TreeSet 本质是 TreeMap 子树的视图，而非副本。对视图的修改会直接反映到原集合。

六、常见问题与设计考量

1. 为何不允许 null 元素？

• 自然排序场景：

  null 无法调用 compareTo 方法，会抛出 NullPointerException。

• 定制排序场景：

  若 Comparator 允许 null（如显式处理），理论上可支持，但 TreeSet 源码未开放此功能（TreeMap 的键不允许为 null）。

2. equals () 与 compareTo () 的一致性

• 若元素的 compareTo 返回 0（视为相等），但 equals 返回 false，TreeSet 会认为元素重复，无法插入。这是因为 TreeMap 通过 compare 判断键是否相等，而非 equals。

3. 性能优化点

• 空间优化：

  使用 PRESENT 占位符避免为每个元素创建独立值对象，减少内存占用。

• 操作委托：

  所有集合操作委托给 TreeMap，复用红黑树成熟的算法实现，降低代码复杂度。

总结：TreeSet 源码设计核心

1. 数据结构选择：

   基于红黑树的 TreeMap 实现，保证有序性和高效的增删查（O (log n)）。

2. 排序逻辑解耦：

   通过 Comparator 或 Comparable 实现排序规则的灵活配置。

3. 视图机制：

   范围查询返回子树视图，操作高效且与原集合联动。

4. 设计约束：

   强制要求元素可比较，牺牲部分灵活性换取排序功能的可靠性。

理解 TreeSet 的源码需重点关注其与 TreeMap 的委托关系，以及红黑树的自平衡机制。

在实际应用中，需根据业务是否需要排序和范围查询，合理选择 TreeSet 或其他集合类（如 HashSet、LinkedHashSet）。

⬅️ TreeSet 🏠 00-Java ➡️ 00-Queue和Deque总览