TreeMap的源码分析(二)【添加元素平衡】

王守钰 2020-03-17 12:03:15

TreeMap的put方法

public V put(K key, V value) {
    // 定义root节点
    Entry<K,V> t = root;
    // 判断root节点是否为空
    if (t == null) {
        compare(key, key); // type (and possibly null) check
        // 创建root节点
        root = new Entry<>(key, value, null);
        // 元素个数为1
        size = 1;
        // 操作次数加一
        modCount++;
        return null;
    }
    // 比较结果
    int cmp;
    // 父节点
    Entry<K,V> parent;
    // split comparator and comparable paths
    // 获取当前比较器
    Comparator<? super K> cpr = comparator;
    // 判断当前比较器是否为空
    if (cpr != null) {
        do {
            // 定义父节点为root节点
            parent = t;
            // 比较父节点和当前要插入节点
            cmp = cpr.compare(key, t.key);
            // 插入节点key小于父节点,进行左查找
            if (cmp < 0)
                t = t.left;
            // 插入节点key大于父节点,进行右查找
            else if (cmp > 0)
                t = t.right;
            // key值相同设置value,返回老值
            else
                return t.setValue(value);
            // 循环遍历
        } while (t != null);
    }
    // 没有默认的比较器
    else {
        // 判断插入的key是否为空,为空抛出异常
        if (key == null)
            throw new NullPointerException();
        @SuppressWarnings("unchecked")
            // 定义key
            Comparable<? super K> k = (Comparable<? super K>) key;
        do {
            // 定义父节点为root节点
            parent = t;
            // 比较父节点和当前要插入的节点
            cmp = k.compareTo(t.key);
            // 插入节点key小于父节点,进行左查找
            if (cmp < 0)
                t = t.left;
            // 插入节点key大于父节点,进行右查找
            else if (cmp > 0)
                t = t.right;
            // key值相同设置value,返回老值
            else
                return t.setValue(value);
            // 循环遍历
        } while (t != null);
    }
    // 定义当前key
    Entry<K,V> e = new Entry<>(key, value, parent);
    // 判断要插入的节点是左节点还是右节点
    if (cmp < 0)
        parent.left = e;
    else
        parent.right = e;
    // 插入后进行平衡操作
    fixAfterInsertion(e);
    // 元素个数加1
    size++;
    // 操作次数加1
    modCount++;
    // 如果插入了新节点返回null
    return null;
}

插入再平衡

插入节点默认是红色节点。

  1. 如果插入节点为root节点,直接将其图为黑色,不用平衡;
  2. 插入元素的父节点如果为黑色,不用平衡;
  3. 插入元素的父节点如果为红色,需要进行平衡:
    • 如果父节点是祖父节点的左节点。
      • 父节点和叔叔节点都为红色,需要将父节点和叔叔节点都标为黑色,祖父节点标记为红色,将祖父节点当做为新的节点,进行下一次循环判断。
      • 父节点为红色,叔叔节点为黑色,插入节点为父节点的左节点。以父节点为圆心节点祖父为旋转节点进行右旋。
      • 父节点为红色,叔叔节点为黑色,插入节点为父节点的右节点。以父节点为圆心节点当前节点为旋转节点进行左旋。然后这样就会变成和上一操作步骤一样了,重复上一个操作方法进行操作。
    • 如果父节点是祖父节点的右节点。
      • 父节点和叔叔节点都为红色,需要将父节点和叔叔节点都标为黑色,祖父节点标记为红色,将祖父节点当做为新的节点,进行下一次循环判断。
      • 父节点为红色,叔叔节点为黑色,插入节点为父节点的右节点。以父节点为圆心节点祖父为旋转节点进行左旋。
      • 父节点为红色,叔叔节点为黑色,插入节点为父节点的左节点。以父节点为圆心节点当前节点为旋转节点进行右旋。然后这样就会变成和上一操作步骤一样了,重复上一个操作方法进行操作。
private void fixAfterInsertion(Entry<K,V> x) {
    // 设置插入节点x为红色节点
    x.color = RED;
    // 当前节点不为空,不是root节点,并且当前节点的父节点是红色
    while (x != null && x != root && x.parent.color == RED) {
        // 当前节点的父节点是否是祖父节点的左子节点
        if (parentOf(x) == leftOf(parentOf(parentOf(x)))) {
            // 定义y为叔叔节点
            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;
}

父节点为祖父节点的左节点插入

插入左子节点
image
右旋
image
变色
image

插入右子节点
image
左旋
image
右旋
image
变色
image

父节点为祖父节点的右节点插入

插入右子节点
image
左旋
image
变色
image
插入左子节点
image
右旋
image
左旋
image
变色
image