TreeSet源码分析

TreeSet源码分析

功能：将Set中的元素按照一定的规则进行排序存储。

在其源码的内部实现中(如下)，可以看到TreeSet时借助了TreeMap来实现的。

    public TreeSet() {
        this(new TreeMap<E,Object>());
    }

    TreeSet(NavigableMap<E,Object> m) {
        this.m = m;
    }

TreeSet.add方法

    public boolean add(E e) {
            //    private static final Object PRESENT = new Object();
        return m.put(e, PRESENT)==null;
    }

从add方法中可以看到直接将e其作为key，value = new Object() 保存在了TreeMap中。

下面来看一下TreeMap类中put方法的内部实现哈

    public TreeMap() {
        comparator = null;
    }

TreeMap.put(key,value)的实现思想比较简单：

1）就是实现了一个二叉排序树：要么是空树，要么满足以下条件：若左子树不空，则左子树所有结点的值均小于根结点的值，若右子树不空，右子树所有结点的值均大于根结点的值；左子树和右子树也是一颗二叉排序树。

2）具体为：从root节点开始遍历，利用比较器comparator来比较node.key与key的大小来确定此key应该存放的位置。注意：comparator可能为null，如果为null，则使用key的natural ordering(自然顺序)，例如：没有指定comparator的String 类型key 的结果就是字典排序。

       public V put(K key, V value) {
        Entry<K,V> t = root;
        if (t == null) {
            compare(key, key); // type (and possibly null) check

            root = new Entry<>(key, value, null);
            size = 1;
            modCount++;
            return null;
        }
        int cmp;
        Entry<K,V> parent;
        // split comparator and comparable paths
        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)
                throw new NullPointerException();
            @SuppressWarnings("unchecked")
                Comparable<? super K> k = (Comparable<? super K>) key;
            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++;
        modCount++;
        return null;
    } 

TreeSet for-each方法的原理

经常我们会这里来遍历已经排好序的TreeSet

        for (String key : treeSet) {
           //...
        }

那么，里面是如何来实现从最小的元素开始依次开始取元素的呢？

简单来说：由于存储是使用的二叉树排序的思想，因此对二叉树进行中序遍历即可得到有序的结果集。

从源码来分析，具体这就要看TreeMap中的keySet()方法的内部实现了

    public Set<K> keySet() {
        return navigableKeySet();
    }

    public NavigableSet<K> navigableKeySet() {
        KeySet<K> nks = navigableKeySet;
        return (nks != null) ? nks : (navigableKeySet = new KeySet<>(this));
    }

这里的KeySet这个类就是关键了。

    static final class KeySet<E> extends AbstractSet<E> implements NavigableSet<E> {
        private final NavigableMap<E, ?> m;
        KeySet(NavigableMap<E,?> map) { m = map; }

        public Iterator<E> iterator() {
            if (m instanceof TreeMap)
                return ((TreeMap<E,?>)m).keyIterator();
            else
                return ((TreeMap.NavigableSubMap<E,?>)m).keyIterator();
        }
    //...省略了不需要关注的代码    
    }               

从上面KeySet的iterator的实现中可以看到，通过如下的方法得到了类型为KeyIterator的迭代器。

        Iterator<K> keyIterator() {
            return new KeyIterator(getFirstEntry());
        }

        /**
         * Returns the first Entry in the TreeMap (according to the TreeMap's
         * key-sort function).  Returns null if the TreeMap is empty.
         */
        final Entry<K,V> getFirstEntry() {
            Entry<K,V> p = root;
            if (p != null)
                while (p.left != null)
                    p = p.left;
            return p;
        } 

getFirstEntry函数的功能：找到排序结果中最小的那个元素。

下面主要看下 KeyIterator 这个类

这个类的next方法就是用来获取TreeSet中的下一个排好序的元素的。

具体实现思路为：由于上面讲解的put方法中我们知道是根据“左根右”的思想来存储的[最小值,中间值,最大值]，因此这里的获取最小值的思路也是如此。

     final class KeyIterator extends PrivateEntryIterator<K> {
        KeyIterator(Entry<K,V> first) {
            super(first);
        }
        public K next() {
            return nextEntry().key;
        }
    }

        final Entry<K,V> nextEntry() {
            Entry<K,V> e = next;
            if (e == null)
                throw new NoSuchElementException();
            if (modCount != expectedModCount)
                throw new ConcurrentModificationException();
            next = successor(e);
            lastReturned = e;
            return e;
        }                  

    /**
     * Returns the successor of the specified Entry, or null if no such.
     */
    static <K,V> TreeMap.Entry<K,V> successor(Entry<K,V> t) {
        if (t == null)
            return null;
        else if (t.right != null) {//判断t节点的右边还有没有元素，如果有取出最小的。
            Entry<K,V> p = t.right;
            while (p.left != null)
                p = p.left;
            return p;
        } else {//如果此时t的右边没有元素可遍历了，则遍历t的父节点的父节点
            Entry<K,V> p = t.parent;
            Entry<K,V> ch = t;
            while (p != null && ch == p.right) {
                ch = p;
                p = p.parent;
            }
            return p;
        }
    }

小结

比较简单哈，就是根据顺序存储，然后按照一定的顺序取出。