按键排序的Map-TreeMap源码分析

总结

总结放前面(这篇挺短的)
1. TreeMap基于红黑树实现,可以对Map中的key排序
2. 它的排序和定位需要依赖比较器或实现 Comparable 接口,也因此不需要key重写hashCode方法和equals方法,就可以排除掉重复的key。
3. TreeMap 的查询、插入、删除效率均没有 HashMap 高,一般只有要对 key 排序时才使用TreeMap。

导言

之前有写过弱平衡的二叉查找树-红黑树,然而很多具体实现都没有提,这篇文章就来谈谈JDK基于红黑树实现的可以按键排序的Map——TreeMap

还是按我的习惯,直接从重要方法读起:

重要方法

put

public V put(K key, V value) {
    //可以猜到Entry应该就是红黑树的节点类
    Entry<K,V> t = root;
    //如果根节点为空
    if (t == null) {
        //这里的compare只是类型检查和null值检查key
        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;
    // 获取比较器
    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);
    }
    // 如果没有比较器 即为实现了Comparable接口,类似的逻辑
    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;
}

看完这个函数的结论:
1. Entry是结点的类
2. 按正常的搜索树方法插入节点之后,调用fixAfterInsertion(e)来调整树。

那么我们分别来看对应的源码:

Entry

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

其实似曾相识,因为之前在HashMap源码分析中看过HashMap用红黑树实现的部分也实现了红黑树的节点,不过当时也没有谈具体方法的实现。

fixAfterInsertion

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

get

public V get(Object key) {
    Entry<K,V> p = getEntry(key);
    return (p==null ? null : p.value);
}

get其实没什么好看的,就是搜索树的搜索。

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

remove

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.
    // 如果p有两个孩子节点
    if (p.left != null && p.right != null) {
        // 获取后继结点
        Entry<K,V> s = successor(p);
        // 用后继节点替换当前节点(颜色不变)
        p.key = s.key;
        p.value = s.value;
        // 然后把p指向s,和后面p只有一个或0个孩子节点的方法通用
        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;
    //没有孩子节点,且要删除的是黑色节点,那么还要通过fixAfterDeletion再继续调整
    } 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;
        }
    }
}
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是由红黑树实现的,不存在什么扩容,构造方法也很简单:

public TreeMap() {
    comparator = null;
}

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

原创文章,作者:彭晨涛,如若转载,请注明出处:https://www.codetool.top/article/%e6%8c%89%e9%94%ae%e6%8e%92%e5%ba%8f%e7%9a%84map-treemap%e6%ba%90%e7%a0%81%e5%88%86%e6%9e%90/