TreeSet源码分析

1、概述

TreeSet 是 Set 接口的另一个实现类，它内部采用自平衡的排序二叉树来存储元素，这种结构保证 TreeSet 集合中没有重复的元素，并且对元素进行排序。

TreeSet 集合存入一个元素时，都会将存入的元素与其他元素比较，这就要求 TreeSet 集合存入的元素必须是实现Comparable接口的compareTo方法的，否则运行时将报异常。

附 Type hierarchy：

①平衡二叉树

平衡二叉树全称平衡查找二叉树，平衡二叉树的每个节点最多有两个子节点，每个节点及其子节点组成的树称为子树。左侧的节点称为“左子树”；右侧的节点称为“右子树”，左子树与右子树的高度差的绝对值不超过1。平衡二叉树结构图如图所示：

2、源码分析

public class TreeSet<E> extends AbstractSet<E> 
    implements NavigableSet<E>, Cloneable, java.io.Serializable 
    {
    // 备份的NavigableMap对象m
    private transient NavigableMap<E, Object> m;

    // 用于反复填充m的value的临时值PRESENT
    private static final Object PRESENT = new Object();

    // 创建TreeSet对象时，若要传入额外的Navigable对象，
    TreeSet(NavigableMap<E, Object> m) {
        this.m = m;
    }

    // 无参构造函数，直接创建一个TreeMap对象
    public TreeSet() {
        this(new TreeMap<E, Object>());
    }

    // 构造函数，传入一个Comparator类型的比较器对象
    public TreeSet(Comparator<? super E> comparator) {
        this(new TreeMap<>(comparator));
    }

    // 构造函数，传入一个集合对象
    public TreeSet(Collection<? extends E> c) {
        this();
        addAll(c);
    }

    // 构造函数，传入一个SortedSet类型的对象，并通过构造函数得到其比较器
    public TreeSet(SortedSet<E> s) {
        this(s.comparator());
        addAll(s);
    }

    // 得到迭代器
    public Iterator<E> iterator() {
        return m.navigableKeySet().iterator();
    }

    // 得到降序的迭代器
    public Iterator<E> descendingIterator() {
        return m.descendingKeySet().iterator();
    }

    public NavigableSet<E> descendingSet() {
        return new TreeSet<>(m.descendingMap());
    }

    // 得到集合的大小
    public int size() {
        return m.size();
    }

    // 判断集合是否为空
    public boolean isEmpty() {
        return m.isEmpty();
    }

    // 使用内部的TreeMap（Backing map）的key来判断HashSet中是否存在重复的元素
    public boolean contains(Object o) {
        return m.containsKey(o);
    }

    // HashSet添加元素，实际上是添加到TreeMap的key中，TreeMap的value为PRESENT常量
    public boolean add(E e) {
        return m.put(e, PRESENT) == null;
    }

    // 删除元素
    public boolean remove(Object o) {
        return m.remove(o) == PRESENT;
    }

    // 清除所有元素
    public void clear() {
        m.clear();
    }

    // 添加集合
    public boolean addAll(Collection<? extends E> c) {
        // Use linear-time version if applicable
        if (m.size() == 0 && c.size() > 0 && c instanceof SortedSet && m instanceof TreeMap) {
            SortedSet<? extends E> set = (SortedSet<? extends E>) c;
            TreeMap<E, Object> map = (TreeMap<E, Object>) m;
            Comparator<? super E> cc = (Comparator<? super E>) set.comparator();
            Comparator<? super E> mc = map.comparator();
            if (cc == mc || (cc != null && cc.equals(mc))) {
                map.addAllForTreeSet(set, PRESENT);
                return true;
            }
        }
        return super.addAll(c);
    }

    // 获取NavigableSet类型的子集合，子集合为TreeSet类型
    public NavigableSet<E> subSet(E fromElement, boolean fromInclusive, E toElement, boolean toInclusive) {
        return new TreeSet<>(m.subMap(fromElement, fromInclusive, toElement, toInclusive));
    }

    // 获取NavigableSet类型的头集合，头集合为TreeSet类型
    public NavigableSet<E> headSet(E toElement, boolean inclusive) {
        return new TreeSet<>(m.headMap(toElement, inclusive));
    }

    // 获取NavigableSet类型的尾集合，尾集合为TreeSet类型
    public NavigableSet<E> tailSet(E fromElement, boolean inclusive) {
        return new TreeSet<>(m.tailMap(fromElement, inclusive));
    }

    // 获取SortedSet类型的子集合
    public SortedSet<E> subSet(E fromElement, E toElement) {
        return subSet(fromElement, true, toElement, false);
    }

    // 获取SortedSet类型的头集合
    public SortedSet<E> headSet(E toElement) {
        return headSet(toElement, false);
    }

    // 获取SortedSet类型的尾集合
    public SortedSet<E> tailSet(E fromElement) {
        return tailSet(fromElement, true);
    }

    // 获取TreeMap的比较器
    public Comparator<? super E> comparator() {
        return m.comparator();
    }

    // 获取第一个元素
    public E first() {
        return m.firstKey();
    }

    // 获取最后一个元素
    public E last() {
        return m.lastKey();
    }

    // 略
    public E lower(E e) {
        return m.lowerKey(e);
    }

    /**
     * @throws ClassCastException
     *             {@inheritDoc}
     * @throws NullPointerException
     *             if the specified element is null and this set uses natural
     *             ordering, or its comparator does not permit null elements
     * @since 1.6
     */
    public E floor(E e) {
        return m.floorKey(e);
    }

    /**
     * @throws ClassCastException
     *             {@inheritDoc}
     * @throws NullPointerException
     *             if the specified element is null and this set uses natural
     *             ordering, or its comparator does not permit null elements
     * @since 1.6
     */
    public E ceiling(E e) {
        return m.ceilingKey(e);
    }

    /**
     * @throws ClassCastException
     *             {@inheritDoc}
     * @throws NullPointerException
     *             if the specified element is null and this set uses natural
     *             ordering, or its comparator does not permit null elements
     * @since 1.6
     */
    public E higher(E e) {
        return m.higherKey(e);
    }

    /**
     * @since 1.6
     */
    public E pollFirst() {
        Map.Entry<E, ?> e = m.pollFirstEntry();
        return (e == null) ? null : e.getKey();
    }

    /**
     * @since 1.6
     */
    public E pollLast() {
        Map.Entry<E, ?> e = m.pollLastEntry();
        return (e == null) ? null : e.getKey();
    }

    // 克隆对象
    public Object clone() {
        TreeSet<E> clone = null;
        try {
            clone = (TreeSet<E>) super.clone();
        } catch (CloneNotSupportedException e) {
            throw new InternalError();
        }

        clone.m = new TreeMap<>(m);
        return clone;
    }

    // 写入ObjectOutputStream序列化对象
    private void writeObject(java.io.ObjectOutputStream s) throws java.io.IOException {
        // Write out any hidden stuff
        s.defaultWriteObject();

        // Write out Comparator
        s.writeObject(m.comparator());

        // Write out size
        s.writeInt(m.size());

        // Write out all elements in the proper order.
        for (E e : m.keySet())
            s.writeObject(e);
    }

    // 从ObjectInputStream流里读取序列化对象
    private void readObject(java.io.ObjectInputStream s) throws java.io.IOException, ClassNotFoundException {
        // Read in any hidden stuff
        s.defaultReadObject();

        // Read in Comparator
        Comparator<? super E> c = (Comparator<? super E>) s.readObject();

        // Create backing TreeMap
        TreeMap<E, Object> tm;
        if (c == null)
            tm = new TreeMap<>();
        else
            tm = new TreeMap<>(c);
        m = tm;

        // Read in size
        int size = s.readInt();

        tm.readTreeSet(size, s, PRESENT);
    }

    private static final long serialVersionUID = -2479143000061671589L;
}

总结：TreeSet本质上是使用TreeMap的key存储元素的。

3、实例

import java.util.TreeSet;

public class Q{
    public static void main(String[] args) {
        TreeSet<Person> ts = new TreeSet<Person>();

        ts.add(new Person("zhangsan", 19));
        ts.add(new Person("lisi", 24));
        ts.add(new Person("wangwu", 27));

        for (Person person : ts) {
            System.out.println(person);
        }
    }
}

class Person implements Comparable<Person>{
    private String name;
    private int age;

    public String getName() {
        return name;
    }
    public void setName(String name) {
        this.name = name;
    }
    public int getAge() {
        return age;
    }
    public void setAge(int age) {
        this.age = age;
    }

    public Person(String name, int age) {
        this.name = name;
        this.age = age;
    }

    @Override
    public int compareTo(Person p) {
        //按照降序排列
        if (this.getAge() - p.getAge() > 0) {
            return -1;
        }
        if (this.getAge() == p.getAge()) {
            return -this.getName().compareTo(p.getName());
        }
        return 1;
    }

    @Override
    public String toString() {
        return "姓名：" + this.getName() + "|" + " 年龄：" + this.getAge();
    }
}

运行结果：

姓名：wangwu| 年龄：27
姓名：lisi| 年龄：24
姓名：zhangsan| 年龄：19