【Java源码】TreeMap类

上一篇分析了HashMap的源码，今天有人在群里面试了阿里说问到了HashMap和TreeMap的比较，所以干脆也来看看TreeMap的实现。

HashMap是Map的非同步实现，而TreeMap也不是同步的，TreeMap和HashMap的最大区别在于元素的放置:

——》HashMap使用哈希表的方式存储元素，所以插入、删除、查询定位比较快

——》而TreeMap则使用RB-Tree（红黑树）的方式来存储元素，这样插入进TreeMap中的元素是有序排列的

一、数据结构

和HashMap一样，TreeMap的内部数据结构也使用实现了一个Map.Entry键值对，但不同的是HashMap中entry是链表的实现，而TreeMap中是红黑树的实现。

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

        /**
         * Make a new cell with given key, value, and parent, and with
         * <tt>null</tt> child links, and BLACK color.
         */
        Entry(K key, V value, Entry<K,V> parent) {
            this.key = key;
            this.value = value;
            this.parent = parent;
        }

二、TreeMap的属性

1、元素排序使用的比较器Comparetor

/**
     * The comparator used to maintain order in this tree map, or
     * null if it uses the natural ordering of its keys.
     *
     * @serial
     */
    private final Comparator<? super K> comparator;

2、红黑树的根节点，初始为空

private transient Entry<K,V> root = null;

3、TreeMap中元素个数

 /**
     * The number of entries in the tree
     */
    private transient int size = 0;

4、和HashMap一样的修改计数器，但是和HashMap中不同的是，它没有用volatile修饰

/**
     * The number of structural modifications to the tree.
     */
    private transient int modCount = 0;

三、重要方法

1、添加元素put()

整体步骤是：从根节点开始遍历红黑树——》根据Comparator决定要插入的分支——》找到数中插入位置并插入成叶节点——》调整红黑树的平衡性

public V put(K key, V value) {
        Entry<K,V> t = root;
        if (t == null) {//如果根节点为null
	    // TBD:
	    // 5045147: (coll) Adding null to an empty TreeSet should
	    // throw NullPointerException
	    //
	    // compare(key, key); // type check
            root = new Entry<K,V>(key, value, null);//将新节点设为root
            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);//比较插入的key和当前遍历的节点，即潜在父节点
                if (cmp < 0)
                    t = t.left;//若比父节点小，则转向父节点的左支继续遍历
                else if (cmp > 0)
                    t = t.right;
                else
                    return t.setValue(value);//若key和父节点相同，则更新该节点的value
            } while (t != null);
        }
        else {
            if (key == null)
                throw new NullPointerException();
            Comparable<? super K> k = (Comparable<? super K>) key;//若没有比较器，则使用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<K,V>(key, value, parent);
        if (cmp < 0)//找到插入位置后，将新节点添加为叶节点上
            parent.left = e;
        else
            parent.right = e;
        fixAfterInsertion(e);//调整红黑树的平衡性
        size++;
        modCount++;
        return null;
    }

1.1——》fixAfterInsertion()调整红黑树，。。。擦，不太懂红黑树，先路过。。。

/** From CLR */
    private void fixAfterInsertion(Entry<K,V> x) {
        x.color = RED;

        while (x != null && x != root && x.parent.color == RED) {
            if (parentOf(x) == leftOf(parentOf(parentOf(x)))) {
                Entry<K,V> y = rightOf(parentOf(parentOf(x)));
                if (colorOf(y) == RED) {
                    setColor(parentOf(x), BLACK);
                    setColor(y, BLACK);
                    setColor(parentOf(parentOf(x)), RED);
                    x = parentOf(parentOf(x));
                } else {
                    if (x == rightOf(parentOf(x))) {
                        x = parentOf(x);
                        rotateLeft(x);
                    }
                    setColor(parentOf(x), BLACK);
                    setColor(parentOf(parentOf(x)), RED);
                    rotateRight(parentOf(parentOf(x)));
                }
            } else {
                Entry<K,V> y = leftOf(parentOf(parentOf(x)));
                if (colorOf(y) == RED) {
                    setColor(parentOf(x), BLACK);
                    setColor(y, BLACK);
                    setColor(parentOf(parentOf(x)), RED);
                    x = parentOf(parentOf(x));
                } else {
                    if (x == leftOf(parentOf(x))) {
                        x = parentOf(x);
                        rotateRight(x);
                    }
                    setColor(parentOf(x), BLACK);
                    setColor(parentOf(parentOf(x)), RED);
                    rotateLeft(parentOf(parentOf(x)));
                }
            }
        }
        root.color = BLACK;
    }

2、读取元素get()

本质上就是对红黑树，或者更一般的说对二叉搜索树进行查找

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
        if (comparator != null)
            return getEntryUsingComparator(key);
        if (key == null)
            throw new NullPointerException();
	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;
    }

3、删除节点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;
    }

3.1——》deleteEntry，红黑树的修改操作

 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
            replacement.parent = p.parent;
            if (p.parent == null)
                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) { // return if we are the only node.
            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;
            }
        }
    }

3.2——》fixAfterDelete，红黑树的调整

/** From CLR */
    private void fixAfterDeletion(Entry<K,V> x) {
        while (x != root && colorOf(x) == BLACK) {
            if (x == leftOf(parentOf(x))) {
                Entry<K,V> sib = rightOf(parentOf(x));

                if (colorOf(sib) == RED) {
                    setColor(sib, BLACK);
                    setColor(parentOf(x), RED);
                    rotateLeft(parentOf(x));
                    sib = rightOf(parentOf(x));
                }

                if (colorOf(leftOf(sib))  == BLACK &&
                    colorOf(rightOf(sib)) == BLACK) {
                    setColor(sib, RED);
                    x = parentOf(x);
                } else {
                    if (colorOf(rightOf(sib)) == BLACK) {
                        setColor(leftOf(sib), BLACK);
                        setColor(sib, RED);
                        rotateRight(sib);
                        sib = rightOf(parentOf(x));
                    }
                    setColor(sib, colorOf(parentOf(x)));
                    setColor(parentOf(x), BLACK);
                    setColor(rightOf(sib), BLACK);
                    rotateLeft(parentOf(x));
                    x = root;
                }
            } else { // symmetric
                Entry<K,V> sib = leftOf(parentOf(x));

                if (colorOf(sib) == RED) {
                    setColor(sib, BLACK);
                    setColor(parentOf(x), RED);
                    rotateRight(parentOf(x));
                    sib = leftOf(parentOf(x));
                }

                if (colorOf(rightOf(sib)) == BLACK &&
                    colorOf(leftOf(sib)) == BLACK) {
                    setColor(sib, RED);
                    x = parentOf(x);
                } else {
                    if (colorOf(leftOf(sib)) == BLACK) {
                        setColor(rightOf(sib), BLACK);
                        setColor(sib, RED);
                        rotateLeft(sib);
                        sib = leftOf(parentOf(x));
                    }
                    setColor(sib, colorOf(parentOf(x)));
                    setColor(parentOf(x), BLACK);
                    setColor(leftOf(sib), BLACK);
                    rotateRight(parentOf(x));
                    x = root;
                }
            }
        }

        setColor(x, BLACK);
    }

总结：和HashMap相比，TreeMap保证了元素的有序性，但牺牲了在查找、删除、添加时的速度，因此TreeMap使用了红黑树作为内置结构，能够尽量弥补速度；所以平时若不是要保证顺序，还是要使用HashMap