TreeMap 源码分析 - litter-fish/ReadSource GitHub Wiki

类继承关系

Tree节点

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

总有的属性

// 比较器
private final Comparator<? super K> comparator;
// 树的根节点
private transient Entry<K,V> root;

// 键值对的数量
private transient int size = 0;

// 修改次数，用于快速失败
private transient int modCount = 0;

构造方法

// comparator用键的顺序做比较
public TreeMap() {
    comparator = null;
}

// 构造方法，提供比较器，用指定比较器排序
public TreeMap(Comparator<? super K> comparator) {
    this.comparator = comparator;
}

// 将Map元素插入TreeMap中，按照键顺序比较
public TreeMap(Map<? extends K, ? extends V> m) {
    comparator = null;
    putAll(m);
}

// 构造方法，指定的参数为SortedMap
// 采用m的比较器排序
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原理

红黑树性质：

节点的颜色只能是红色和黑色
根节点一定是黑色
所有叶子节点的子节点都是空节点且都是黑色的
不存在两个连续的红色节点
任何节点到子树上的黑色节点数量相同

插入情况：情形1：如果插入的节点既是根节点，此时将颜色设置成黑色即可

情形2：如果插入节点的父节点是黑色的，插入的话并不会改变红黑树的性质

情形3：如果插入节点的父节点P和叔父节点U都为红色

处理：交换祖父节点和父亲、叔父节点的颜色，此时祖父节点颜色为红色，

情形4：如果插入节点的父节点P是红色，而叔父节点U是黑色，且新增节点N是P的右孩子，P是G的左孩子

处理：父节点进行一次左旋操作，转换为情形5进行处理

情形5：如果插入节点的父节点P是红色，而叔父节点U是黑色或缺失，同时新增节点N是P的左孩子且P也是G的左孩子

处理：

交换父节点和祖父节点颜色
将祖父节点做右旋操作

插入源码

public V put(K key, V value) {
    Entry<K,V> t = root; // 树的根节点
    // 如果树是一颗空树
    if (t == null) {
        compare(key, key); // type (and possibly null) check
        // 将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) {
        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); // 如果插入key值已经存在，直接替换
        } while (t != null); // 遍历树找到正确的插入位置，红黑树是一个二叉搜索树，so，比根小的值往左走，比根大的往右走
    }
    // 没有指定比较器，使用自然排序方式
    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;
}

红黑树的调整

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

            // 情景3：父节点P和叔父节点U都为红色，
            // 需要重新绘制P、U为黑色，祖父G为红色，但是G节点可能是根节点也可能其父节点也是红色，因此需要按照情形1递归处理
            // 如果叔父节点是红色的，父节点也是红色的，且插入节点也是红色的
            if (colorOf(y) == RED) {

                setColor(parentOf(x), BLACK); // 设置父节点为黑色
                setColor(y, BLACK); // 设置叔父节点为黑色
                setColor(parentOf(parentOf(x)), RED); // 将祖父节点设置为红色

                // 祖父节点是红色的，此时将祖父节点这颗子树当做一个整体进行循环处理。
                // 此时如果祖父节点既是根节点则，不满足循环条件，跳出
                // 如果不是根节点则将其当做一个红色节点整体进行循环处理。
                x = parentOf(parentOf(x));
            } else { // 如果叔父节点是黑色的，此时父节点红色，叔父节点红色
                // 情形4.1：如果父节点P是红色，而叔父节点U是黑色，且新增节点N是P的右孩子，P是G的左孩子，此种情况，需要将节点P进行一次左旋操作，然后按照情形5.1进行操作
                if (x == rightOf(parentOf(x))) {
                    // 获取父节点，并进行一次左旋
                    x = parentOf(x);
                    rotateLeft(x);
                }
                // 情形5.1：如果父节点P是红色，且其叔父节点U是黑色或缺失，同时新增节点N是P的左孩子且P也是G的左孩子，
                // 此时将节点祖父节点G进行一次右旋，并替换G和P节点的颜色
                setColor(parentOf(x), BLACK);
                setColor(parentOf(parentOf(x)), RED);
                // 祖父节点右旋
                rotateRight(parentOf(parentOf(x)));
            }
        } else {
            Entry<K,V> y = leftOf(parentOf(parentOf(x)));
            // 情景3：父节点P和叔父节点U都为红色
            if (colorOf(y) == RED) {
                setColor(parentOf(x), BLACK);
                setColor(y, BLACK);
                setColor(parentOf(parentOf(x)), RED);
                x = parentOf(parentOf(x));
            } else {
                // 情形4.2：如果父节点P是红色，而叔父节点U是黑色，且新增节点N是P的左孩子，P是G的右孩子，此种情况，需要将节点P进行一次右旋操作，然后按照情形5.2进行操作
                if (x == leftOf(parentOf(x))) {
                    x = parentOf(x);
                    // 父节点进行右旋
                    rotateRight(x);
                }
                // 情形5.2：如果父节点P是红色，且其叔父节点U是黑色或缺失，同时新增节点N是P的右孩子且P也是G的右孩子，此时将节点祖父节点G进行一次左旋，并替换G和P节点的颜色
                setColor(parentOf(x), BLACK);
                setColor(parentOf(parentOf(x)), RED);
                // 以祖父节点做左旋
                rotateLeft(parentOf(parentOf(x)));
            }
        }
    }
    root.color = BLACK;
}

左旋转

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“上浮”，因此将p的父节点指向新的节点R的父节点， 节点P“下沉”
        r.parent = p.parent;
        // 如果父节点为null，表示该节点是一个根节点
        if (p.parent == null)
            root = r;
        // 如果旋转的节点是父节点的左孩子，则将父节点的左孩子指向调整后的节点
        else if (p.parent.left == p)
            p.parent.left = r;
        else
            p.parent.right = r;
        // 将旋转的节点赋值给r节点的左孩子
        r.left = p;
        // 待旋转的节点下沉，其父节点指向旋转节点的右孩子
        p.parent = r;
    }
}

右旋

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节点被上沉，故将其父节点指向待旋转节点的父节点
        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;
    }
}

红黑树的删除原理及TreeMap的remove实现

删除步骤

将红黑树当做一颗二叉查找树，将节点删除
通过旋转、着色重新恢复成一颗红黑树

删除节点

如果删除的是红色的不会影响红黑树
如果删除的节点是黑色，则经过的路径黑色少了一个，违背红黑树原则

删除情况

如果删除的节点没有孩子，直接删除
删除的节点只有一个孩子，直接用孩子顶替
如果删除的节点存在两个孩子，则查找左子树的最大值或右子树的最小值，替换节点，然后转换成删除一个孩子的情况，如下图所示，删除x节点

删除节点x后，存在以下一些情况情况1：被删除节点为根节点或者是红色节点无需操作

情况2：N的兄弟节点sib是红色

处理：

交换兄弟节点和其父节点的颜色
再将父节点左旋
设置父节点的右孩子为新的兄弟节点，此时被转换成情况3-1

情况3：父节点可能是红色也可能是黑色 3-1：父节点红色的情况

处理：设置兄弟节点为红色，同时将删除替换的节点设置成其父类的节点，此时不满足节点是黑色的循环条件，跳出循环，最后将节点设置成黑色

3-2：父节点黑色情况

处理：设置兄弟节点sib为红色，此时经过父节点的黑色节点数相同了，但不经过父节点的黑色节点比其多了一个，此时将经过父节点这颗子树当做一个整体进行循环处理

情况4：删除后替换的N节点是父节点的左子树，兄弟节点的左孩子SL是红色、右孩子SR是黑色的。

处理：

交换叔父节点和叔父节点的左节点的颜色
再将兄弟节点进行右旋操作
重新设置sib节点将会转换成情况5再做处理

情况5：兄弟节点的右节点是红色，此时左节点可能是红色也可能是黑色的

处理

Sib颜色设置成P的颜色，P和sib的右节点设置成黑色
父节点P左旋

删除源码

public V remove(Object key) {
    // 根据key找到节点
    Entry<K,V> p = getEntry(key);
    if (p == null)
        return null;
    // 获取待删除节点的value值
    V oldValue = p.value;
    // 删除节点
    deleteEntry(p);
    // 返回删除节点的值
    return oldValue;
}

根据key获取节点

final Entry<K,V> getEntry(Object key) {
    // 使用自定义的比较器获取节点
    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;
}

获取右节点

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;
        // 循环获取左节点，直到不存在，此时返回了右节点最小key值得节点
        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;
    }
}