Java从入门到放弃（十四）集合框架之TreeMap源码

我们经常需要对一些集合按照指定的规则进行排序，比如学生按照学号排序，或者按照成绩排序，集合里面有专门排序的集合，如TreeMap。TreeMap里面是使用的红黑树结构。

构造方法

   private final Comparator<? super K> comparator;
   private transient Entry<K,V> root; 
   private transient int size = 0;
   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 | ClassNotFoundException cannotHappen) {
        }
    }

有三个变量，comparator是指定的比较器，root是红黑树的根节点，size是集合内数据的长度。可以看到每一个方法都有指定比较器，在TreeMap中比较器是必须的。因为没有比较器就没有办法对集合内的元素进行对比，也就无法进行排序。

put方法

    public V put(K key, V value) {
        Entry<K,V> t = root;       //根节点
        if (t == null) {
            compare(key, key); // type 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);  //使用比较器比较t和key的大小
                if (cmp < 0)
                    t = t.left;     //如果key小于t，把t更改为t的左子树
                else if (cmp > 0)
                    t = t.right;    //如果key大于t,把t更改为t的右子树
                else
                    return t.setValue(value);    //如果相等，更改t的值
            } while (t != null);    //一直遍历到null，找到合适的父节点
        }
        else {
            if (key == null)
                throw new NullPointerException();   //如果比较器和key都是null，抛出异常
            @SuppressWarnings("unchecked")
                Comparable<? super K> k = (Comparable<? super K>) key; //把key强制转换为comparable
            do {
                parent = t;
                cmp = k.compareTo(t.key);   //此处和上面是一样的，就是循环找出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);    //构建key节点
        if (cmp < 0)
            parent.left = e;    //左子树
        else
            parent.right = e;  //右子树
        fixAfterInsertion(e);   //调整树的结构
        size++;
        modCount++;
        return null;
    }
 final int compare(Object k1, Object k2) {
        return comparator==null ? ((Comparable<? super K>)k1).compareTo((K)k2)
            : comparator.compare((K)k1, (K)k2);
    }

put方法就是从root节点开始查找合适的父节点，节点是由一个内部类Entry构成的，看一下Entry类的实现

   static final class Entry<K,V> implements Map.Entry<K,V> {
        K key;     //存储的key值
        V value;   //存储的value值
        Entry<K,V> left;  //左子树
        Entry<K,V> right;   //右子树
        Entry<K,V> parent;   //父节点
        boolean color = BLACK;   //默认黑色

        Entry(K key, V value, Entry<K,V> parent) {
            this.key = key;    
            this.value = value;
            this.parent = parent;
        }
        public K getKey() {
            return key;
        }
        public V getValue() {
            return value;
        }

        public V setValue(V value) {
            V oldValue = this.value;
            this.value = value;
            return oldValue;
        }

        public boolean equals(Object o) {
            if (!(o instanceof Map.Entry))
                return false;
            Map.Entry<?,?> e = (Map.Entry<?,?>)o;

            return valEquals(key,e.getKey()) && valEquals(value,e.getValue());
        }

        public int hashCode() {
            int keyHash = (key==null ? 0 : key.hashCode());
            int valueHash = (value==null ? 0 : value.hashCode());
            return keyHash ^ valueHash;
        }

        public String toString() {
            return key + "=" + value;
        }
    }

Entry里面包含节点的key-value值，左右节点和父节点。方法也很简单，只有value的set，get方法，key只有get方法，所以key是不可修改的，重写了hashcode和eauals以及tostring方法。

remove和getEntry方法

    public V remove(Object key) {
        Entry<K,V> p = getEntry(key);    //获取对应的节点
        if (p == null)
            return null; 

        V oldValue = p.value;
        deleteEntry(p);    //删除节点
        return oldValue;   //返回旧的值
    }
    final Entry<K,V> getEntry(Object key) {
        // Offload comparator-based version for sake of performance
        if (comparator != null)
            return getEntryUsingComparator(key); //如果指定了比较器，执行这个方法，内部和下面的是一样的
        if (key == null)
            throw new NullPointerException();
        @SuppressWarnings("unchecked")
            Comparable<? super K> k = (Comparable<? super K>) key;  //转换为comparable
        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;
    }

这里也是找到对应的节点，然后用deleteEntry方法进行删除，里面涉及到红黑树的左旋右旋以及颜色的调整，这里不做解释。

clear方法

    public void clear() {
        modCount++;
        size = 0;
        root = null;
    }

这里和之前的集合不一样，之前的ArrayList等都是遍历置为null,这里是直接把root置为null;

代码实例

        TreeMap<String,String> treeMap = new TreeMap();
        treeMap.put("China","Beijing");
        treeMap.put("american","Washington");
        treeMap.put("UnitedKingdom","London");
        treeMap.put("Russia","Moscow");
        System.out.println(treeMap);

输出为：{China=Beijing, Russia=Moscow, UnitedKingdom=London, american=Washington}
这里就是没有指定比较器，但是String实现了comparable接口，根据字典顺序排序，因为ASCII表中大写字母的值小于小写字母的值，所以大写字母和小写比较，小写字母会被认为更大。
如果实体类要进行排序要实现comparable或者compartor接口。具体可参考Java从入门到放弃（四）Comparable 和Comparator排序。