图解TreeMap的红黑树平衡 *** 作fixAfterInsertion()，接着手撕红黑树添加节点

图解TreeMap的红黑树平衡 *** 作fixAfterInsertion()，接着手撕红黑树添加节点 一、前言

啥也不想说，就卷、卷技术；手撕红黑树搞起。

1、红黑树简介

红黑树就是一种平衡的二叉查找树，其有五个特点：

1.每个节点要么是红⾊，要么是⿊⾊；
2. 根节点⼀定是⿊⾊的；
3. 每个叶⼦节点⼀定是⿊⾊的NULL节点；
4. 如果⼀个节点是红⾊，那么它的左右⼦节点⼀定都是⿊⾊的；
5. 从任意⼀个节点到叶⼦节点，所经过的⿊⾊节点的数量⼀样多；

2、TreeMap简介

TreeMap 继承于AbstractMap ，其是一个 有序的key-value集合，内部基于 红黑树 实现；
TreeMap 根据 其key的自然顺序进行排序，或者在构造方法中指定 Comparator 进行排序；
TreeMap的基本 *** 作 containsKey()、get()、put() 和 remove() 的时间复杂度是 log(n)。
另外，TreeMap是非同步 的。 它的迭代器是fail-fast 的。

二、put()方法 1、put *** 作实现原理

1. 如果树为null，则构建一个TreeMapEntry设置为当前的root；
2. 排序方式优先使用自定义比较器，找到要插入节点的父节点，然后插入；
3. 执行fixAfterInsertion(e)方法维护红黑树的平衡；

public V put(K key, V value) {
    Entry t = root;
    if (t == null) {
        compare(key, key); // type (and possibly null) check

        root = new Entry<>(key, value, null);
        size = 1;
        modCount++;
        return null;
    }
    int cmp;
    Entry parent;
    // split comparator and comparable paths
    Comparator 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 {
        if (key == null)
            throw new NullPointerException();
        @SuppressWarnings("unchecked")
            Comparable k = (Comparable) 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 e = new Entry<>(key, value, parent);
    if (cmp < 0)
        parent.left = e;
    else
        parent.right = e;
    fixAfterInsertion(e);
    size++;
    modCount++;
    return null;
}

我们重点看一下红黑树的平衡是如何维护的

2、维护红黑树的平衡 – fixAfterInsertion()

建议大家细细拼，这一块还挺绕！！

新插入的节点颜色默认为红色。

平衡 *** 作流程如下：

1. 当节点不是跟节点且节点的父节点颜色为红色时进行while循环 *** 作。

1. 节点的父节点为爷爷节点的左子节点时；取爷爷节点的右子节点；

• 1）爷爷节点的右子节点颜色为红色时；
  将爷爷节点的颜色改为红色，父节点和父节点的兄弟节点置为黑色；
  
  
• 2）否则：
  1、如果x是其父节点的右子节点：令当前节点为其父节点，接着执行左旋转 *** 作；
  
  
  2、将x的父节点和爷爷节点的颜色对换。然后对爷爷节点进行右旋转；
  
  

2. 节点的父节点为爷爷节点的右子节点时，取爷爷节点的左孩子；

• 1）当爷爷节点的左子节点颜色为红色时；
  父节点和父节点的兄弟节点设置为黑色，爷爷节点设置为红色；
  将爷爷节点作为新插入的节点x，遍历调整；
  
  

• 2）否则：
  1、如果x是其父亲的左孩子：令x为其父节点，然后进行右旋 *** 作；
  
  
  2、将父节点设置为黑色，爷爷节点设置为红色，然后以爷爷节点为旋转点进行左旋转；
  
  

最后将根节点的颜色设置为黑色；

private void fixAfterInsertion(Entry x) {
    x.color = RED;

    while (x != null && x != root && x.parent.color == RED) {
        if (parentOf(x) == leftOf(parentOf(parentOf(x)))) {
            Entry 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 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;
}

图解解决，我们接着手撕个红黑树的添加节点流程。

三、手撕红黑树添加节点

public class RedBlackTree, V> {

    public static void main(String[] args) {
        RedBlackTree rdTree = new RedBlackTree();
        rdTree.add(1, 1);
        rdTree.add(3, 1);
        rdTree.add(2, 1);
        RedBlackTree.TreeNode root = rdTree.root;
        System.out.println(root);
    }

    
    private static final boolean RED = true;
    private static final boolean BLACK = false;

    
    class TreeNode {
        public K key;
        public V value;
        public TreeNode left;
        public TreeNode right;
        public boolean color;

        public TreeNode(K key, V value) {
            this.key = key;
            this.value = value;
            this.left = null;
            this.right = null;
            this.color = RED;
        }
    }

    private TreeNode root;

    private int size;

    public RedBlackTree() {
        root = null;
        size = 0;
    }

    
    public boolean isRed(TreeNode node) {
        if (node == null) {
            return BLACK;
        }
        return node.color;
    }

    public void add(K key, V value) {
        root = add(root, key, value);
        root.color = BLACK;
    }

    public TreeNode add(TreeNode node, K key, V value) {
        if (node == null) {
            size++;
            return new TreeNode(key, value);
        }
        // 1.往树中添加节点
        if (key.compareTo(node.key) < 0) {
            node.left = add(node.left, key, value);
        } else if (key.compareTo(node.key) > 0) {
            node.right = add(node.right, key, value);
        } else {
            node.value = value;
        }
        // 2.维护红黑树的结构
        // 2.1 如果node节点的右子节点为红色的，而左子节点不是红色的，左旋转
        if (isRed(node.right) && !isRed(node.left)) {
            node = leftRotate(node);
        }
        // 2.2 如果node节点的左子节点、左子节点的左子节点是红色的，右旋转
        if (isRed(node.left) && isRed(node.left.left)) {
            node = rightRotate(node);
        }
        // 2.3 如果node节点的左子节点和右子节点都是红节点，则颜色翻转
        if (isRed(node.left) && isRed(node.right) && !isRed(node)) {
            flipColor(node);
        }
        return node;
    }


    
    private TreeNode leftRotate(TreeNode node) {
        TreeNode x = node.right;
        node.right = x.left;
        x.left = node;
        x.color = node.color;
        node.color = RED;
        return x;
    }

    
    private TreeNode rightRotate(TreeNode node) {
        TreeNode x = node.left;
        node.left = x.right;
        x.right = node;
        x.color = node.color;
        node.color = RED;
        return x;
    }

    
    private void flipColor(TreeNode node) {
        node.color = RED;
        node.left.color = BLACK;
        node.right.color = BLACK;
    }
}

只要我们技术足够卷、没有hold不住的技术面，奥利给！