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;

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

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

构造方法

	/**
     * 无比较器,需泛型实现 Comparable 接口
     */
    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) {
        }
    }

增删改查

增

/**
 * 添加
 */
public V put(K key, V value) {
    Entry<K,V> t = root;
    //根节点不存在,
    if (t == null) {
        //检测是否有比较器 或 实现Comparable接口
        compare(key, key); 

        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) {
        //开始迭代,直到节点为空,此时可插入节点,父节点为parent
        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);
    }
    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;
}

恢复红黑树性质

/** From CLR */
private void fixAfterInsertion(Entry<K,V> x) {
    // 默认为 red
    x.color = RED;
    // x 不为 null , x 不为 root, x 父节点为 red 节点(上溢) --> 需进行恢复
    while (x != null && x != root && x.parent.color == RED) {
        // x 的父节点 是 x 的祖父节点的 左节点
        if (parentOf(x) == leftOf(parentOf(parentOf(x)))) {
            // y 是 x 的祖父节点的 右节点, 即 y 是 x 的叔叔节点
            Entry<K,V> y = rightOf(parentOf(parentOf(x)));
            // 叔叔节点是 red 节点  --> 上溢 LL/LR 情况
            if (colorOf(y) == RED) {
                // 父节点染黑
                setColor(parentOf(x), BLACK);
                // 叔叔节点染黑
                setColor(y, BLACK);
                // 祖父节点染红
                setColor(parentOf(parentOf(x)), RED);
                // 将新修改节点继续迭代, 直至平衡
                x = parentOf(parentOf(x));
            } 
            // 叔叔节点是 black 节点
            else {
                // x 是父节点右节点 --> LR 情况
                if (x == rightOf(parentOf(x))) {
                    // 左旋转父节点, x 节点上升一位, 父节点下移一位
                    x = parentOf(x);
                    // 经过旋转, 现在新添加的节点相当于原添加节点的父节点
                    rotateLeft(x);
                }
                // x 是父节点左节点, LR 情况经左旋转到达 LL 情况
                // 父节点染黑
                setColor(parentOf(x), BLACK);
                // 祖父节点染红
                setColor(parentOf(parentOf(x)), RED);
                // 右旋转, 祖父节点成为父节点右节点
                rotateRight(parentOf(parentOf(x)));
            }
        } 
        // x 的父节点 是 x 的祖父节点的 右节点
        else {
            Entry<K,V> y = leftOf(parentOf(parentOf(x)));
            // 叔叔节点是 red 节点  --> 上溢 RR/RL 情况
            if (colorOf(y) == RED) {
                setColor(parentOf(x), BLACK);
                setColor(y, BLACK);
                setColor(parentOf(parentOf(x)), RED);
                x = parentOf(parentOf(x));
            } 
            // 叔叔节点是 black 节点
            else {
                // x 是父节点左节点 --> RL 情况
                if (x == leftOf(parentOf(x))) {
                    x = parentOf(x);
                    rotateRight(x);
                }
                // x 是父节点右节点, RL 情况经右旋转到达 RR 情况
                setColor(parentOf(x), BLACK);
                setColor(parentOf(parentOf(x)), RED);
                rotateLeft(parentOf(parentOf(x)));
            }
        }
    }
    //根节点永远是 black
    root.color = BLACK;
}

左右旋转

/** From CLR */
private void rotateLeft(Entry<K,V> p) {
    if (p != null) {
        Entry<K,V> r = p.right;
        p.right = r.left;
        if (r.left != null)
            r.left.parent = p;
        r.parent = p.parent;
        if (p.parent == null)
            root = r;
        else if (p.parent.left == p)
            p.parent.left = r;
        else
            p.parent.right = r;
        r.left = p;
        p.parent = r;
    }
}

/** From CLR */
private void rotateRight(Entry<K,V> p) {
    if (p != null) {
        Entry<K,V> l = p.left;
        p.left = l.right;
        if (l.right != null) l.right.parent = p;
        l.parent = p.parent;
        if (p.parent == null)
            root = l;
        else if (p.parent.right == p)
            p.parent.right = l;
        else p.parent.left = l;
        l.right = p;
        p.parent = l;
    }
}

比较方法

/**
 * 优先使用比较器
 * 传的是 Object ,需强制转换
 */
@SuppressWarnings("unchecked")
final int compare(Object k1, Object k2) {
    return comparator==null ? ((Comparable<? super K>)k1).compareTo((K)k2)
        : comparator.compare((K)k1, (K)k2);
}

Node类

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

删

/**
 * 门面
 */
public V remove(Object key) {
    //检测是否存在
    Entry<K,V> p = getEntry(key);
    if (p == null)
        return null;

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

/**
 * Delete node p, and then rebalance the tree.
 *
 * 度为 0 --  父节点对应子树设置为 null
 * 度为 1 --  父节点对应子树设置为删除节点子树
 * 度为 2 --  从左右子树中寻找节点替代,并对替代节点进行更新
 */
private void deleteEntry(Entry<K,V> p) {
    modCount++;
    size--;

    // 有左右子树,返回右子树的最左左叶子节点
    if (p.left != null && p.right != null) {
        //寻找后继节点进行替代
        Entry<K,V> s = successor(p);
        //修改要删除节点的 key 和 value
        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;

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

        // 删除的节点是 black 节点
        if (p.color == BLACK)
            // 恢复红黑树平衡
            fixAfterDeletion(replacement);
    } else if (p.parent == null) { // 只有一个节点
        root = null;
    } 
    // 叶子节点,直接更新父节点 左右子树为 null
    else { 
        if (p.color == BLACK)
            // 恢复红黑树平衡 
            fixAfterDeletion(p);

        //p 有父节点
        if (p.parent != null) {
            //p 为父节点的左子树
            if (p == p.parent.left)
                p.parent.left = null;
            //p 为父节点的右子树
            else if (p == p.parent.right)
                p.parent.right = null;
            //gc 回收
            p.parent = null;
        }
    }
}

恢复红黑树性质

// 删除节点是 black 节点 --> 传入的是删除后替代删除节点的节点 或 删除节点是 black 叶子节点
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));
            // 兄弟节点是 red 节点 (black 叶子节点时进行修正)--> 将兄弟节点转化为 black 节点
            if (colorOf(sib) == RED) {
                setColor(sib, BLACK);
                setColor(parentOf(x), RED);
                // 左旋转, 右节点上提
                rotateLeft(parentOf(x));
                // 更新
                sib = rightOf(parentOf(x));
            }
            // 兄弟节点经转化为必然 black 节点
            // 兄弟节点左右子节点是 black, 没有 red 节点借给删除节点
            // colorOf( null ) = BLACK
            if (colorOf(leftOf(sib)) == BLACK && colorOf(rightOf(sib)) == BLACK) {
                setColor(sib, RED);
                // 有可能再次迭代, 进行修复
                x = parentOf(x);
            } else {
                // 兄弟节点有一个 red 节点, 且为 black 右节点, red 左节点 --> 右旋转变为有左节点
                if (colorOf(rightOf(sib)) == BLACK) {
                    setColor(leftOf(sib), BLACK);
                    setColor(sib, RED);
                    rotateRight(sib);
                    sib = rightOf(parentOf(x));
                }
                // 必然 : 兄弟节点有一个 red 节点, 且为 black 左节点, red 右节点 --> 统一操作
                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;
            }
        }
    }
    // x = parentOf(x) 或 x = root
    setColor(x, BLACK);
}

successor - 后继

/**
 * 返回指定Entry的后继；如果没有，则返回null 
 */
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;
    } 
    // t != null && t.right == null
    // 没有右子树, 只能向上找
    else {
        Entry<K,V> p = t.parent;
        Entry<K,V> ch = t;
        //向上进行迭代,寻找后继节点
        /*
        		   5
        	   /       \
        	  3          8
        	/   \       /  \
           1      2    7    10 
        
        * 7 后继节点是 8 ,需要从父节点寻找,左子树迭代上去
        */
        while (p != null && ch == p.right) {
            ch = p;
            p = p.parent;
        }
        return p;
    }
}

改

//和增一样,替换

查

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

/**
 * 核心
 */
final Entry<K,V> getEntry(Object key) {
    // 优先使用比较器
    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;
}

迭代