TreeMap 源码

问题

1. 如何用 TreeMap 实现一致性 Hash ？

概述

• TreeMap 存储K-V键值对，通过红黑树（R-B tree）实现；
• 红黑树结构天然支持排序，默认情况下通过Key值的自然顺序进行排序；

继承结构

• 实现了 SortedMap 接口

  可以按key的大小来遍历

• 实现了 NavigableMap 接口

  可查询一些离目标key最近的元素的方法

TreeMap 使用红黑树的存储结构，时间复杂度为 O(log n)，

红黑树

红黑树规则：

• 节点分为红色或者黑色；
• 根节点必为黑色；
• 叶子节点都为黑色，且为null；
• 连接红色节点的两个子节点都为黑色（红黑树不会出现相邻的红色节点）；
• 从任意节点出发，到其每个叶子节点的路径中包含相同数量的黑色节点；
• 新加入到红黑树的节点为红色节点；

红黑树自平衡基本操作：

• 变色：在不违反上述红黑树规则特点情况下，将红黑树某个node节点颜色由红变黑，或者由黑变红；
• 左旋：逆时针旋转两个节点，让一个节点被其右子节点取代，而该节点成为右子节点的左子节点；
• 右旋：顺时针旋转两个节点，让一个节点被其左子节点取代，而该节点成为左子节点的右子节点；

源码实现

属性

View Code

/**
 * 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;

private transient Entry<K,V> root;

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

• 两种比较器的使用

  1. Key 实现了 Commparator 接口；
  2. 通过构造方法传入比较器；

• 根节点

  树的根节点，所有的树都存储在一棵树；

内部类

红黑树结构。

View Code

/**
 * Node in the Tree.  Doubles as a means to pass key-value pairs back to
 * user (see Map.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;

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

Entry 的实现结构：

构造方法

View Code

/**
 * Constructs a new, empty tree map, using the natural ordering of its
 * keys.  All keys inserted into the map must implement the {@link
 * Comparable} interface.  Furthermore, all such keys must be
 * <em>mutually comparable</em>: {@code k1.compareTo(k2)} must not throw
 * a {@code ClassCastException} for any keys {@code k1} and
 * {@code k2} in the map.  If the user attempts to put a key into the
 * map that violates this constraint (for example, the user attempts to
 * put a string key into a map whose keys are integers), the
 * {@code put(Object key, Object value)} call will throw a
 * {@code ClassCastException}.
 */
public TreeMap() {
    comparator = null;
}

/**
 * Constructs a new, empty tree map, ordered according to the given
 * comparator.  All keys inserted into the map must be <em>mutually
 * comparable</em> by the given comparator: {@code comparator.compare(k1,
 * k2)} must not throw a {@code ClassCastException} for any keys
 * {@code k1} and {@code k2} in the map.  If the user attempts to put
 * a key into the map that violates this constraint, the {@code put(Object
 * key, Object value)} call will throw a
 * {@code ClassCastException}.
 *
 * @param comparator the comparator that will be used to order this map.
 *        If {@code null}, the {@linkplain Comparable natural
 *        ordering} of the keys will be used.
 */
public TreeMap(Comparator<? super K> comparator) {
    this.comparator = comparator;
}

/**
 * Constructs a new tree map containing the same mappings as the given
 * map, ordered according to the <em>natural ordering</em> of its keys.
 * All keys inserted into the new map must implement the {@link
 * Comparable} interface.  Furthermore, all such keys must be
 * <em>mutually comparable</em>: {@code k1.compareTo(k2)} must not throw
 * a {@code ClassCastException} for any keys {@code k1} and
 * {@code k2} in the map.  This method runs in n*log(n) time.
 *
 * @param  m the map whose mappings are to be placed in this map
 * @throws ClassCastException if the keys in m are not {@link Comparable},
 *         or are not mutually comparable
 * @throws NullPointerException if the specified map is null
 */
public TreeMap(Map<? extends K, ? extends V> m) {
    comparator = null;
    putAll(m);
}

/**
 * Constructs a new tree map containing the same mappings and
 * using the same ordering as the specified sorted map.  This
 * method runs in linear time.
 *
 * @param  m the sorted map whose mappings are to be placed in this map,
 *         and whose comparator is to be used to sort this map
 * @throws NullPointerException if the specified map is null
 */
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) {
    }
}

两种比较器

• 传入定义的 Comparator 比较器；
• Key 实现 Comparable 接口；

get(Object key)

View Code

/**
 * Returns the value to which the specified key is mapped,
 * or {@code null} if this map contains no mapping for the key.
 *
 * <p>More formally, if this map contains a mapping from a key
 * {@code k} to a value {@code v} such that {@code key} compares
 * equal to {@code k} according to the map's ordering, then this
 * method returns {@code v}; otherwise it returns {@code null}.
 * (There can be at most one such mapping.)
 *
 * <p>A return value of {@code null} does not <em>necessarily</em>
 * indicate that the map contains no mapping for the key; it's also
 * possible that the map explicitly maps the key to {@code null}.
 * The {@link #containsKey containsKey} operation may be used to
 * distinguish these two cases.
 *
 * @throws ClassCastException if the specified key cannot be compared
 *         with the keys currently in the map
 * @throws NullPointerException if the specified key is null
 *         and this map uses natural ordering, or its comparator
 *         does not permit null keys
 */
public V get(Object key) {
    Entry<K,V> p = getEntry(key);
    return (p==null ? null : p.value);
}

/**
 * Returns this map's entry for the given key, or {@code null} if the map
 * does not contain an entry for the key.
 *
 * @return this map's entry for the given key, or {@code null} if the map
 *         does not contain an entry for the key
 * @throws ClassCastException if the specified key cannot be compared
 *         with the keys currently in the map
 * @throws NullPointerException if the specified key is null
 *         and this map uses natural ordering, or its comparator
 *         does not permit null keys
 */
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();

    // 将key强转为Comparable
    @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)
            // 小于0，从左子树查找
            p = p.left;
        else if (cmp > 0)
            // 大于0，从右子树查找
            p = p.right;
        else
            // 如果相等，则返回
            return p;
    }
    return null;
}

/**
 * Version of getEntry using comparator. Split off from getEntry
 * for performance. (This is not worth doing for most methods,
 * that are less dependent on comparator performance, but is
 * worthwhile here.)
 */
final Entry<K,V> getEntryUsingComparator(Object key) {
    @SuppressWarnings("unchecked")
        K k = (K) key;
    Comparator<? super K> cpr = comparator;
    if (cpr != null) {
        // 从根元素开始遍历
        Entry<K,V> p = root;
        while (p != null) {
            int cmp = cpr.compare(k, p.key);
            if (cmp < 0)
                // 小于0，从左子树查找
                p = p.left;
            else if (cmp > 0)
                // 大于0，从右子树查找
                p = p.right;
            else
                // 相等，则返回
                return p;
        }
    }
    return null;
}

二分查找的思想：

• 从root遍历整个树；
• 如果待查找的 key 比当前遍历的 key 小，则在其左子树中查找；
• 如果待查找的 key 比当前遍历的 key 大，则在其右子树中查找；
• 如果待查找的 key 与当前遍历的 key 相等，则找到了该元素，直接返回；

put(K key, V value)

View Code

/**
 * Associates the specified value with the specified key in this map.
 * If the map previously contained a mapping for the key, the old
 * value is replaced.
 *
 * @param key key with which the specified value is to be associated
 * @param value value to be associated with the specified key
 *
 * @return the previous value associated with {@code key}, or
 *         {@code null} if there was no mapping for {@code key}.
 *         (A {@code null} return can also indicate that the map
 *         previously associated {@code null} with {@code key}.)
 * @throws ClassCastException if the specified key cannot be compared
 *         with the keys currently in the map
 * @throws NullPointerException if the specified key is null
 *         and this map uses natural ordering, or its comparator
 *         does not permit null keys
 */
public V put(K key, V value) {
    Entry<K,V> t = root;
    if (t == null) {
        // 没有根节点，直接插入根节点
        compare(key, key); // type (and possibly null) check
        
        root = new Entry<>(key, value, null);
        size = 1;
        modCount++;
        return null;
    }

    // key 比较的结果
    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);
            if (cmp < 0)
                // 左子树查找
                t = t.left;
            else if (cmp > 0)
                // 右子树查找
                t = t.right;
            else
                // 节点已存在并替换
                return t.setValue(value);
        } while (t != null);
    }
    else {
        // 使用 key 实现的比较器
        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;
}

• 插入的元素如果是根节点，则直接涂成黑色即可，不用平衡；
• 插入的元素的父节点如果为黑色，不需要平衡；
• 插入的元素的父节点如果为红色，则做需要平衡，源码通过 fixAfterInsertion(e) 方法进行自平衡；

fixAfterInsertion(Entry<K,V> x)

插入后再平衡。

View Code

/** 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)))) {
            // a）如果父节点是祖父节点的左节点
            // y为叔叔节点
            Entry<K,V> y = rightOf(parentOf(parentOf(x)));
            if (colorOf(y) == RED) {
                // 情况1）如果叔叔节点为红色  
                // （1）将父节点设为黑色
                setColor(parentOf(x), BLACK);
                // （2）将叔叔节点设为黑色
                setColor(y, BLACK);
                // （3）将祖父节点设为红色
                setColor(parentOf(parentOf(x)), RED);
                // （4）将祖父节点设为新的当前节点
                x = parentOf(parentOf(x));
            } else {
                // 如果叔叔节点为黑色
                // 情况2）如果当前节点为其父节点的右节点
                if (x == rightOf(parentOf(x))) {
                    // （1）将父节点设为当前节点
                    x = parentOf(x);
                    // （2）以新当前节点左旋
                    rotateLeft(x);
                }
                // 情况3）如果当前节点为其父节点的左节点（如果是情况2）则左旋之后新当前节点正好为其父节点的左节点了）
                // （1）将父节点设为黑色
                setColor(parentOf(x), BLACK);
                // （2）将祖父节点设为红色
                setColor(parentOf(parentOf(x)), RED);
                // （3）以祖父节点为支点进行右旋
                rotateRight(parentOf(parentOf(x)));
            }
        } else {
            // b）如果父节点是祖父节点的右节点
            // y是叔叔节点
            Entry<K,V> y = leftOf(parentOf(parentOf(x)));
            if (colorOf(y) == RED) {
                // 情况1）如果叔叔节点为红色
                // （1）将父节点设为黑色
                setColor(parentOf(x), BLACK);
                // （2）将叔叔节点设为黑色
                setColor(y, BLACK);
                // （3）将祖父节点设为红色
                setColor(parentOf(parentOf(x)), RED);
                // （4）将祖父节点设为新的当前节点
                x = parentOf(parentOf(x));
            } else {
                // 如果叔叔节点为黑色
                // 情况2）如果当前节点为其父节点的左节点
                if (x == leftOf(parentOf(x))) {
                    // （1）将父节点设为当前节点
                    x = parentOf(x);
                    // （2）以新当前节点右旋
                    rotateRight(x);
                }

                // 情况3）如果当前节点为其父节点的右节点（如果是情况2）则右旋之后新当前节点正好为其父节点的右节点了）
                // （1）将父节点设为黑色
                setColor(parentOf(x), BLACK);
                // （2）将祖父节点设为红色
                setColor(parentOf(parentOf(x)), RED);
                // （3）以祖父节点为支点进行左旋
                rotateLeft(parentOf(parentOf(x)));
            }
        }
    }
    // 黑色节点，不必再做旋转
    root.color = BLACK;
}

fixAfterInsertion(Entry<K,V> x) 进行自平衡处理，以下是红黑树处理逻辑：

场景	无需调整	只需【变色】	【旋转 + 变色】
场景1	父节点为黑色时，插入子节点（默认插入节点为红色）	空树插入根节点，将根节点红色变为黑色	父节点为红色左节点，叔父节点为黑色，插入左子节点，那么通过【左左节点旋转】
场景2	-	父节点和叔父节点都为红色	父节点为红色左节点，叔父节点为黑色，插入右子节点，那么通过【左右节点旋转】
场景3	-	-	父节点为红色右节点，叔父节点为黑色，插入左子节点，那么通过【右左节点旋转】
场景4	-	-	父节点为红色右节点，叔父节点为黑色，插入右子节点，那么通过【右右节点旋转】

红黑树参考文档

如果符合再平衡的条件 (父节点为红色)，分六种情况处理：

父节点为祖父节点的左节点

• 父节点为红色，叔父节点也为红色（场景一）

  1. 将父节点设为黑色
  2. 将叔父节点设为黑色
  3. 将祖父节点设为红色
  4. 将祖父节点设为新的当前节点

• 父节点为红色，叔父节点也为黑色，且当前节点是其父节点的右节点（场景二）

  1. 将父节点设为当前节点
  2. 以新当前节点左旋
  3. 此时进入（场景三）

• 父节点为红色，叔父节点也为黑色，且当前节点是其父节点的左节点（场景三）

  1. 将父节点设为黑色
  2. 将祖父节点设为红色
  3. 以祖父节点为支点进行右旋

父节点为祖父节点的右节点

• 父节点为红色，叔父节点也为红色（场景四）

  1. 将父节点设为黑色
  2. 将叔父节点设为黑色
  3. 将祖父节点设为红色
  4. 将祖父节点设为新的当前节点

• 父节点为红色，叔父节点也为黑色，且当前节点是其父节点的右节点（场景五）

  1. 将父节点设为当前节点
  2. 以新当前节点右旋
  3. 此时进入（场景六）

• 父节点为红色，叔父节点也为黑色，且当前节点是其父节点的左节点（场景六）

  1. 将父节点设为黑色
  2. 将祖父节点设为红色
  3. 以祖父节点为支点进行左旋

红黑树动画演示

remove(Object key)

View Code

/**
 * Removes the mapping for this key from this TreeMap if present.
 *
 * @param  key key for which mapping should be removed
 * @return the previous value associated with {@code key}, or
 *         {@code null} if there was no mapping for {@code key}.
 *         (A {@code null} return can also indicate that the map
 *         previously associated {@code null} with {@code key}.)
 * @throws ClassCastException if the specified key cannot be compared
 *         with the keys currently in the map
 * @throws NullPointerException if the specified key is null
 *         and this map uses natural ordering, or its comparator
 *         does not permit null keys
 */
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.
 */
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
        // 把替换节点直接放到当前节点的位置上（相当于删除了p，并把替换节点移动过来了）
        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;
        }
    }
}

• 如果删除的位置有两个叶子节点，则从其右子树中取最小的元素放到删除的位置，然后把删除位置移到替代元素的位置，进入下一步。
• 如果删除的位置只有一个叶子节点（有可能是经过第一步转换后的删除位置），则把那个叶子节点作为替代元素，放到删除的位置，然后把这个叶子节点删除。
• 如果删除的位置没有叶子节点，则直接把这个删除位置的元素删除。
• 针对红黑树，如果删除位置是黑色节点，还需要做再平衡。
• 如果有替代元素，则以替代元素作为当前节点进入再平衡。
• 如果没有替代元素，则以删除的位置的元素作为当前节点进入再平衡，平衡之后再删除这个节点。

fixAfterDeletion(Entry<K,V> x)

删除元素后再平衡。

View Code

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

• 真正删除的肯定是黑色节点才会进入到再平衡阶段。
• 因为删除的是黑色节点，导致整颗树不平衡了，所以这里我们假设把删除的黑色赋予当前节点，这样当前节点除了它自已的颜色还多了一个黑色；

  1. 如果当前节点是根节点，则直接设置为黑色，不需要再平衡；
  2. 如果当前节点是红节点，则直接设置为黑色，不需要平衡；
  3. 如果当前节点是黑节点，
  4. 如果当前节点是黑节点，则我们要通过旋转把这个多出来的黑色不断的向上传递到一个红色节点

红黑树动画演示

总结

• TreeMap 的存储结构是一棵红黑树；
• TreeMap 中的元素是有序的，按 key 的顺序排列；
• TreeMap 可以按照范围查找元素；

小知识

左旋

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

左旋：

逆时针旋转两个节点，让一个节点被其右子节点取代，而该节点成为右子节点的左子节点

左旋操作步骤如下：

首先断开节点PL与右子节点G的关系，同时将其右子节点的引用指向节点C2；然后断开节点G与左子节点C2的关系，同时将G的左子节点的应用指向节点PL

右旋

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

右旋：

顺时针旋转两个节点，让一个节点被其左子节点取代，而该节点成为左子节点的右子节点

右旋操作步骤如下：

首先断开节点G与左子节点PL的关系，同时将其左子节点的引用指向节点C2；然后断开节点PL与右子节点C2的关系，同时将PL的右子节点的应用指向节点G

数据结构可视化
红黑树原理