---
title: "04-TreeSet源码分析"
created: 2025-12-02
tags:
  - Java
---
# TreeSet源码分析

> [!note] 定位说明
> Set 系列收尾深读篇：TreeSet 是 TreeMap 的壳——add 就是 put(key, 一个全局共享的占位对象)。体会"组合优于继承"的设计。

### **一、类定义与核心属性**

#### **1. 类声明**

```java
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. 无参构造（自然排序）**

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

#### **2. 传入 Comparator 的构造函数**

```java
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)`**：元素插入**

```java
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)`**：元素查询**

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

- **核心逻辑**：

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

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

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

- **核心逻辑**：

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

#### **4.** `iterator()`**：迭代器**

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

- **迭代顺序**：

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

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

```java
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` 方法确定元素位置：

```java
// 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. 存储关系**

| **TreeSet 方法** | **对应 TreeMap 操作** |
| --- | --- |
| `add(e)` | `put(e, PRESENT)` |
| `remove(e)` | `remove(e)` |
| `contains(e)` | `containsKey(e)` |
| `size()` | `size()` |
| `first()` | `firstKey()` |
| `last()` | `lastKey()` |

#### **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`）。
---
⬅️ [[03-TreeSet|TreeSet]] 🏠 [[00-Java|00-Java]] ➡️ [[00-Queue和Deque总览|00-Queue和Deque总览]]