【源码剖析】TreeMap和TreeSet

TreeMap

概念特点

• 基于红黑树实现的有序key-value集合，对元素自动排序(比较器)
• 排序、无重复、重复值覆盖
• 查找、插入、删除时间复杂度O(logn)
• 线程不安全（线程安全问题使用ConcurrentSkipListMap）

构造方法

四个构造方法：

• 无参构造，将比较器comparator设置为null,按key自然顺序排序
• 带比较器的构造方法，需要自己实现比较器
• 构造包含指定map集合（无序）的元素，并且使用自然顺序进行插入
• 带sortedMap的构造函数，sortedMap是有序的，使用buildFromSorted()方法将sortedMap集合中的元素插入到TreeMap中

    public TreeMap() {
        comparator = null;
    }

    public TreeMap(Comparator<? super K> comparator) {
        this.comparator = comparator;
    }


    public TreeMap(Map<? extends K, ? extends V> m) {
        comparator = null;
        putAll(m);
    }
   
    public TreeMap(SortedMap<K, ? extends V> m) {
        comparator = m.comparator();
        try {
            buildFromSorted(m.size(), m.entrySet().iterator(), null, null);
        } catch (java.io.IOException cannotHappen) {
        } catch (ClassNotFoundException cannotHappen) {
        }
    }
    

内部组成

内部变量值。

private final Comparator<? super K> comparator;
private transient Entry<K,V> root;//root指向树的根节点
private transient int size = 0;//当前键值对的个数
private transient int modCount = 0;//记录树结构调整的次数，用来实现Fail-Fast 机制，如果在迭代这些集合的过程中，有其他线程修改了这些集合，就会抛出ConcurrentModificationException异常
//红黑树常量
private static final boolean RED   = false;
private static final boolean BLACK = true;

内部数结点的实现类：

static final class Entry<K,V> implements Map.Entry<K,V> {
    K key;
    V value;
    Entry<K,V> left;
    Entry<K,V> right;
    Entry<K,V> parent;//对于根节点，其父节点为 null
    boolean color = BLACK;
    Entry(K key, V value, Entry<K,V> parent) {
        this.key = key;
        this.value = value;
        this.parent = parent;
    }
}

方法执行

保存键值对 put

注：如果初始化TreeMap构造函数时，没有传递Comparator类，是不允许插入key为nul的键值对的；如果实现了Comparator，则可以传递key为nul的键值对。

• 如果是第一次保存，则新建一个节点，root指向它
• 如果不是，则进行遍历，如果找到节点值，就更改，找不到则去寻找它的位置，parent指向父节点，插入节点。
• 调整红黑树。

public V put(K key, V value) {
    Entry<K,V> t = root;
    /**
    * 如果根节点为空，则新建一个结点，设置 root 指向它，即新添加的元素为根节点
    * 并且设置集合的大小 size=1，以及 modCount+1
    */
    if (t == null) {
        compare(key, key); // type (and possibly null) check

        root = new Entry<>(key, value, null);
        size = 1;
        modCount++;//检测结构性变化
        return null;
    }
    /**
    * 如果不是第一次添加元素，则执行以下代码，添加的关键步骤是找寻父节点。
    * 寻找父节点根据是否设置了 comparator 分为两种情况
    */
    int cmp;
    Entry<K,V> parent;
    // split comparator and comparable paths
    Comparator<? super K> cpr = comparator;
    //如果设置了 comparator 进行以下操作
    if (cpr != null) {
    	//通过 do-while 循环不断遍历树，调用比较器对key进行比较
        do {
            parent = t;
            cmp = cpr.compare(key, t.key);
            if (cmp < 0)
                t = t.left;
            else if (cmp > 0)
                t = t.right;
            else
                //遇到key相等，直接将新值覆盖到原值上
                return t.setValue(value);
        } while (t != null);
        //当退出循环时，parent就指向待插结点的父节点
    }
    //没有 comparator 的情况（新插入的元素，按照key的自然排序）
    else {
    	//如果 key=null，直接抛出异常
		//注意：上面构造TreeMap传入了Comparator，是可以允许key为null的
        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;
}

根据键获取值

//就是根据key找到对应的节点p，找到节点后获取值p.value
public V get(Object key) {
    Entry<K,V> p = getEntry(key);
    return (p==null ? null : p.value);
}

final Entry<K,V> getEntry(Object key) {
    // Offload comparator-based version for sake of performance
    //如果 comparator 不为null，调用单独的getEntryUsingComparator方法
    if (comparator != null)
        return getEntryUsingComparator(key);
        
    //如果 comparator 为 null（key实现了Comparable接口，使用该接口的compareTo方法进行比较）    
    //如果 key 为null，抛出异常
    if (key == null)
        throw new NullPointerException();
    //如果 key 不为 null
    @SuppressWarnings("unchecked")
    Comparable<? super K> k = (Comparable<? super K>) key;
    Entry<K,V> p = root;
    while (p != null) {
        int cmp = k.compareTo(p.key);
        if (cmp < 0)
            p = p.left;
        else if (cmp > 0)
            p = p.right;
        else
            return p;
    }
    return null;
}    

查看是否包含某个值

TreeMap可以高效按键进行查找，但如果根据值查找则需要遍历。

• 查找后继算法：
  • 如果当前节点有右孩子，则后继节点为右子树中最小的节点
  • 如果当前节点没有右孩子，则后继节点为节点的某祖先节点。从当前节点开始向上找，如果它是父节点的右孩子，则继续找父节点，直到它不是右孩子或者父节点为空，第一个非右孩子节点的父节点就是后继节点，如果父节点为空，则后继节点为null。

public boolean containsValue(Object value) {
    for (Entry<K,V> e = getFirstEntry(); e != null; e = successor(e))
        //valEquals就是比较值，逐个进行比较
        if (valEquals(value, e.value))
            return true;
    return false;
}

//getFirstEntry 方法返回第一个节点
//第一个节点就是最左边的节点
final Entry<K,V> getFirstEntry() {
    Entry<K,V> p = root;
    if (p != null)
        while (p.left != null)
            p = p.left;
    return p;
}

//successor方法返回给定节点的后继
static <K,V> TreeMap.Entry<K,V> successor(Entry<K,V> t) {
    if (t == null)
        return null;
    else if (t.right != null) {//如果当前节点的右子树不为空
        Entry<K,V> p = t.right;
        while (p.left != null)
            p = p.left;
        return p;
    } else {//如果当前节点的右子树为空
        Entry<K,V> p = t.parent;
        Entry<K,V> ch = t;
        while (p != null && ch == p.right) {
            ch = p;
            p = p.parent;
        }
        return p;
    }
}

根据键删除键值对

删除一个节点时，根据位置不同有不同的删除方法，三种情况：

1. 叶子节点：直接修改父节点对应的应用为null即可
2. 只有一个孩子：在父节点和孩子节点之间直接建立连接
3. 与两个孩子：先找到后继节点，找到后，替换当前节点的内容为后继节点，然后再删除后继节点。因为这个后继节点一定没有做孩子，所有就将两个孩子的情况转换成了上面两种情况了。

//根据key找到节点，调用 deleteEntry 删除节点，返回原来的值
public V remove(Object key) {
    Entry<K,V> p = getEntry(key);
    if (p == null)
        return null;

    V oldValue = p.value;
    deleteEntry(p);
    return oldValue;
}
private void deleteEntry(Entry<K,V> p) {
    modCount++;
    size--;

    // If strictly internal, copy successor's element to p and then make p
    // point to successor.
    //如果左右孩子均不为空
    //		获取后继结点
    //		替换当前节点的内容为后继节点
    //		转换成一个孩子或者叶子节点的情况
    if (p.left != null && p.right != null) {
        Entry<K,V> s = successor(p);
        p.key = s.key;
        p.value = s.value;
        p = s;
    } // p has 2 children

	//有一个孩子的情况
    // Start fixup at replacement node, if it exists.
    Entry<K,V> replacement = (p.left != null ? p.left : p.right);
    if (replacement != null) {
        // Link replacement to parent 
        //在p的父节点和replacement之间建立链接
        replacement.parent = p.parent;
        if (p.parent == null)//p的父节点为空，则修改root指向新的根
            root = replacement;
        else if (p == p.parent.left)
            p.parent.left  = replacement;
        else
            p.parent.right = replacement;

        // Null out links so they are OK to use by fixAfterDeletion.
        p.left = p.right = p.parent = null;

        // Fix replacement 重新平衡树
        if (p.color == BLACK)
            fixAfterDeletion(replacement);
    } else if (p.parent == null) { // //如果树中只有一个根节点
        root = null;
    } else { //  No children. Use self as phantom replacement and unlink.
        if (p.color == BLACK)
            fixAfterDeletion(p);
        if (p.parent != null) {
            if (p == p.parent.left)
                p.parent.left = null;
            else if (p == p.parent.right)
                p.parent.right = null;
            p.parent = null;
        }
    }
}

TreeMap总结

1. 实现Map接口，内部红黑树。
2. 按键有序、要求实现Comparable接口或提供一个Comparator对象
3. 根据键保存、查找、删除效率高，为O(h),h为树的高度

TreeSet

概念特点

• 底层通过HashMap实现
• 排序，无重复，排重
• 线程不安全

构造方法

//Constructs a set backed by the specified navigable map
TreeSet(NavigableMap<E,Object> m) {
    this.m = m;
}
/**
Constructs a new, empty tree set, sorted according to the natural ordering of its elements. 
All elements inserted into the set must implement the {@link Comparable} interface.
*/
public TreeSet() {
    this(new TreeMap<E,Object>());
}
/**
Constructs a new, empty tree set, sorted according to the specified comparator.  
All elements inserted into the set must be <i>mutually comparable</i> by the specified comparator
*/
public TreeSet(Comparator<? super E> comparator) {
    this(new TreeMap<>(comparator));
}    
public TreeSet(Collection<? extends E> c) {
    this();
    addAll(c);
}    
public TreeSet(SortedSet<E> s) {
    this(s.comparator());
    addAll(s);
}    

内部组成

内部使用NavigableMap接口类型，实现类还是TreeMap
value使用一个Object类型的常量PRESENT 代替。

//The backing map
private transient NavigableMap<E,Object> m;//TreeMap实现了NavigableMap接口
// Dummy value to associate with an Object in the backing Map
private static final Object PRESENT = new Object();

方法执行

添加元素

public boolean add(E e) {
    return m.put(e, PRESENT)==null;
}

查看是否包含元素

public boolean contains(Object o) {
    return m.containsKey(o);
}

删除元素

public boolean remove(Object o) {
    return m.remove(o)==PRESENT;
}

总结

实现Set接口，有序，内部实现基于TreeMap

1. 有序，不重复（排重）
2. 添加、删除元素、判断元素是否存在，效率高
3. 要求实现Comparable接口或者通过构造方法提供Comparator对象