TreeSet 源码解析

简介

TreeSet 是 Set 接口的有序实现，底层包装一棵 TreeMap（红黑树），因此元素按「自然顺序」或定制 Comparator 保持有序，并支持范围查询。

• 实现 NavigableSet（间接 SortedSet），提供 first / last / lower / higher / subSet 等。
• 不是线程安全的；并发可用 ConcurrentSkipListSet。
• 不允许「无法比较」的元素（自然排序下传入 null 会抛 NullPointerException）。
• 去重依赖 Comparable / Comparator（相等即去重），查找 / 插入 / 删除复杂度 O(log n)。

TreeSet 的所有元素实际上是 TreeMap 的 key，value 统一用占位对象 PRESENT，与 HashSet→HashMap、LinkedHashSet→LinkedHashMap 的设计如出一辙。

底层数据结构分析

TreeSet 把元素当作 TreeMap 的 key 存入，从而复用红黑树的有序性与范围查询能力。

// 承载元素的 NavigableMap（默认是 TreeMap）
private transient NavigableMap<E,Object> m;

// 统一的占位 value
private static final Object PRESENT = new Object();

构造时即创建 TreeMap：

public TreeSet() {
    this(new TreeMap<E,Object>()); // 默认基于 TreeMap
}

核心成员变量

private transient NavigableMap<E,Object> m;       // 底层红黑树（TreeMap）
private static final Object PRESENT = new Object(); // 占位 value

源码分析

增删查（委托给 TreeMap）

public boolean add(E e) {
    return m.put(e, PRESENT) == null; // 用 TreeMap 的红黑树保证有序与去重
}
public boolean remove(Object o) {
    return m.remove(o) == PRESENT;
}
public boolean contains(Object o) {
    return m.containsKey(o);
}

范围查询（Set 的核心价值）

public E first() { return m.firstKey(); }   // 最小元素
public E last()  { return m.lastKey(); }    // 最大元素
public E lower(E e) { return m.lowerKey(e); }  // 严格小于 e 的最大元素
public NavigableSet<E> subSet(E from, boolean fromInclusive,
                              E to, boolean toInclusive) {
    return new TreeSet<>(m.subMap(from, fromInclusive, to, toInclusive));
}

与 HashSet / LinkedHashSet 的对比

维度	HashSet	LinkedHashSet	TreeSet
底层	HashMap	LinkedHashMap	TreeMap
顺序	无序	插入顺序	key 有序
范围查询	不支持	不支持	支持
允许 null	允许	允许	否（自然排序下）
复杂度	平均 O(1)	平均 O(1)	O(log n)

经典 Demo

import java.util.TreeSet;

public class TreeSetDemo {
    public static void main(String[] args) {
        TreeSet<Integer> set = new TreeSet<>();
        set.add(3); set.add(1); set.add(2);
        System.out.println(set);                        // [1, 2, 3] 自动升序
        System.out.println(set.subSet(1, true, 3, false)); // [1, 2]
        System.out.println(set.descendingSet());           // [3, 2, 1]
    }
}