天天看点

TreeMap 源码解析

之前了解了ArrayList、LinkedList、HashMap、ConcurrentHashMap,接下来再来了解下TreeMap。TreeMap 底层直接维护了一个红黑树,按照惯例,先从继承 及构造函数开始。

public class TreeMap<K,V> extends AbstractMap<K,V> implements NavigableMap<K,V>, Cloneable, java.io.Serializable
{
    // 比较器
    private final Comparator<? super K> comparator;

    // 跟节点
    private transient Entry<K,V> root;

    // 树中实体(节点)的个数
    private transient int size = 0;

    // 对树修改的次数
    private transient int modCount = 0;

    // 无参构造
    public TreeMap() {
        comparator = null;
    }

    // 比较器构造
    public TreeMap(Comparator<? super K> comparator) {
        this.comparator = comparator;
    }

    // Map 集合构造
    public TreeMap(Map<? extends K, ? extends V> m) {
        comparator = null;
        putAll(m);
    }

    // 有序Map构造
    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) {
        }
    }
}
           

构造函数及全局变量已经添加注释,也是非常容易理解的,接下来以无参狗造为基础,通过添加元素,一步一步解析TreeMap。

一、put() 函数

先上源码,通过源码解释更清楚明白(先屏蔽无用的代码):

public V put(K key, V value) {
        Entry<K,V> t = root;
        if (t == null) {
            compare(key, key); // type (and possibly null) check
            // 第一次添加(跟节点为 null),将新增的节点设置为跟节点
            root = new Entry<>(key, value, null);
            size = 1;
            modCount++;
            return null;
        }
        ...
    }
           

从上述源码可以清楚的看到,当初始化 TreeMap 后的第一次 put 操作时,root 为 null ,直接把第一个 put 的数据设置为跟节点。接下来继续添加第二个元素,源码如下:

public V put(K key, V value) {
        Entry<K,V> t = root;
        if (t == null) {
           ...
        }
        int cmp;
        Entry<K,V> parent;
        // split comparator and comparable paths
        Comparator<? super K> cpr = comparator;
        if (cpr != 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;
    }
           

第二次添加元素时,跟节点不为空,所以之前的代码不会再执行。接下来可以看到会对 key 进行一个非空的判断,说明在 TreeMap 中 key 不能为 null 。在下面的 do while 中,从跟节点开始 依次和 本次添加的元素的 key 进行比较,从而找到存放该数据的节点位置。接下来,同样根据默认的排序规则进行比较,将新创建的节点放到指定的位置上。由 TreeMap 的静态内部类可知,新创建的节点 默认为 黑色,源码如下:

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;
        boolean color = BLACK;

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

由红黑树的特性可知,每次新增加的节点均为黑色是显然不对的,所以在添加新的节点之后必须要对红黑树做一定的调整,使其成为真正的红黑树,而调整红黑树的方法就是 fixAfterInsertion(e) 。可以看到它是由新添加的节点开始一步步调整的,源码如下:

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;
    }
           

这个就不一步步梳理了,因为整体来看还是比较简单的,无非就是通过一系列的判断进行左旋或者右旋使其符合红黑树的特性,想要了解红黑树 可以参考另一篇博文。

二、remove(K key); 方法

public V remove(Object key) {
        Entry<K,V> p = getEntry(key);
        if (p == null)
            return null;

        V oldValue = p.value;
        deleteEntry(p);
        return oldValue;
    }
           

看起来也是非常简单,通过key 获取对应的节点,判断节点是否为空。而真正的删除在 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

        // 获取删除节点的左节点或者右节点
        Entry<K,V> replacement = (p.left != null ? p.left : p.right);
    
        if (replacement != null) {
            
            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;

            // 左右父节点都置为空 以便fixAfterDeletion 使用.
            p.left = p.right = p.parent = null;

            // Fix replacement
            if (p.color == BLACK)
                fixAfterDeletion(replacement);
        } else if (p.parent == null) { // 删除的节点 没有子节点 也没有父节点,说明删除的是跟节点,直接将 root 设置为 null 。
            root = null;
        } else { //  没有子节点
            
            if (p.color == BLACK)
                fixAfterDeletion(p);
            // 如果有父节点,删除节点为左节点 或者 右节点 直接设置为null
            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;
            }
        }
    }
           

删除一个节点,首先判断该节点有没有左右节点,如果有,则获取右节点的最左下角的那个节点并返回,源码如下。

/**
     * 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) {
            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;
        }
    }
           

删除节点之后 红黑树的转变,源码如下:

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);
    }
           

这个 删除之后的转变和添加新节点的转变是非常类似的,可以参照一个红黑树进行解读。TreeMap 的简单了解就到这里,更深层次的解读,以后继续分享。

继续阅读