Java实现红黑树的插入【数据结构】

250 阅读5分钟

本笔记的图和脑图都来自大佬的博客,感谢,感谢!!!

1、红黑树的特性

  1. 每个节点或者是黑色,或者是红色
  2. 根节点是黑色
  3. 每个叶子节点(NIL)是黑色。【注意:这里是叶子节点,是指为空(NIL或NULL)的叶子节点!】
  4. 如果一个节点是红色,则它的子节点必须是黑色的。不可以同时存在两个红色节点相连。
  5. 从一个节点到该节点的子孙节点的所有路径上包含相同数目的黑节点。如果一个阶段存在黑色子节点,那么该节点肯定有两个子节点。

image.png

2、插入操作

  1. 查找插入的位置:插入节点的颜色必须是红色,插入黑色节点会破坏平衡。
  2. 插入后自平衡

image.png

红黑树插入情境分析

情景1:红黑树为空树

  1. 将插入节点作为根节点,并将根节点设置为黑色

情景2:插入点的key已存在

  1. 替换原节点的value值

情景3:插入节点的父节点是黑节点

  1. 直接插入,不会影响自平衡

情景4:插入节点的父节点是红节点

根据性质2:父节点为红色,一定不是根节点;根据性质4:爷节点肯定是黑色。

情景4.1 叔叔节点存在,且为红色

image.png

  1. 将父节点和叔节点改为黑色
  2. 将爷节点改为红色
  3. 递归处理爷节点
情景4.2 叔节点不存在或者为黑节点,且插入节点的父节点是爷节点的左孩子

image.png

情景4.2.1 插入节点是父节点的左孩子(LL双红情况)

image.png

  1. 将父节点设置为黑色,将爷节点设置为红色
  2. 对爷节点进行右旋。
情景4.2.2 插入节点是父节点的右孩子(LR双红情况)

image.png

  1. 对父节点进行左旋
  2. 将父节点设置当前节点,得到LL双红情况
  3. 将父节点设置黑色,将爷节点设置红色
  4. 对爷节点进行右旋
情景4.3 叔节点不存在或为黑节点,且插入节点的父节点是爷节点的右孩子

针对这种情况,需要把4.2的逆过来就可以了,不再赘述。

3、代码实现

3.1、RBTree

package p0930;

/**
 * @author: jdq8576
 * @date: 2021-09-30 18:52
 **/
public class RBTree<K extends Comparable<K>, V> {

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

    /**
     * 树根的引用
     *
     * @param node
     * @return
     */
    private RBNode root;

    public RBNode getRoot() {
        return root;
    }

    private RBNode parentOf(RBNode node) {
        if (node != null) {
            return node.parent;
        }
        return null;
    }

    /**
     * 节点是否为红色
     */

    private boolean isRed(RBNode node) {
        if (node != null) {
            return node.color == RED;
        }
        return false;
    }

    /**
     * 是否为黑色
     */

    private boolean isBlack(RBNode node) {
        if (node != null) {
            return node.color == BLACK;
        }
        return false;
    }

    /**
     * 设置节点为红色
     */
    private void setRed(RBNode node) {
        if (node != null) {
            node.color = RED;
        }
    }


    /**
     * 设置节点为黑色
     */
    private void setBlack(RBNode node) {
        if (node != null) {
            node.color = BLACK;
        }
    }

    /**
     * 中序打印二叉树
     */
    public void inOrderPrint() {
        inOrderPrint(root);
    }

    private void inOrderPrint(RBNode node) {
        if (node != null) {
            inOrderPrint(node.left);
            System.out.println("key:" + node.key + ",value:" + node.value);
            inOrderPrint(node.right);
        }
    }

    private void leftRotate(RBNode node) {
        RBNode nR = node.right;
        node.right = nR.left;
        if (nR.left != null) {
            nR.left.parent = node;
        }
        if (node.parent != null) {
            nR.parent = node.parent;
            if (node == node.parent.left) {
                node.parent.left = nR;
            } else {
                node.parent.right = nR;
            }
        } else {
            this.root = nR;
            this.root.parent = null;
        }

        node.parent = nR;
        nR.left = node;
    }

    private void rightRotate(RBNode node) {
        RBNode nL = node.left;
        node.left = nL.right;
        if (nL.right != null) {
            nL.right.parent = node;
        }
        if (node.parent != null) {
            nL.parent = node.parent;
            if (node.parent.left == node) {
                node.parent.left = nL;
            } else {
                node.parent.right = nL;
            }
        } else {
            this.root = nL;
            this.root.parent = null;
        }

        node.parent = nL;
        nL.right = node;
    }

    /**
     * 公开的插入方法
     */

    public void insert(K key, V value) {
        RBNode node = new RBNode();
        node.setKey(key);
        node.setValue(value);
        // 新节点一定要是红色
        node.setColor(RED);
        insert(node);
    }

    private void insert(RBNode node) {
        // 查找node的父节点

        RBNode parent = null;

        RBNode x = this.root;
        while (x != null) {
            parent = x;
            // cmp > 0 说明node.key 大于 x.key 
            // cmp == 0  两者相等 直接替换
            // cmp < 0 node.key 小于 x.key
            int cmp = node.key.compareTo(x.key);
            if (cmp > 0) {
                x = x.right;
            } else if (cmp < 0) {
                x = x.left;
            } else {
                x.setValue(node.value);
                return;
            }
        }
        node.parent = parent;
        if (parent != null) {
            if (parent.key.compareTo(node.key) > 0) {
                parent.left = node;
            } else {
                parent.right = node;
            }
        } else {
            this.root = node;
        }

        // 需要调用修复红黑树平衡的方法
        insertFixUp(node);

    }

    /**
     * 插入后修复红黑树平衡的方法
     */
    private void insertFixUp(RBNode node) {
        this.root.setColor(BLACK);
        RBNode parent = parentOf(node);
        RBNode gparent = parentOf(parent);
        if (parent != null && isRed(parent)) {
            // 如果父节点是红色 那么一定存在爷爷节点,因为根节点不能为红色
            RBNode uncle = null;
            if (parent == gparent.left) {
                uncle = gparent.right;
                if (uncle != null && isRed(uncle)) {
                    setBlack(parent);
                    setBlack(uncle);
                    setRed(gparent);
                    insertFixUp(gparent);
                    return;
                }
                if (uncle == null || isBlack(uncle)) {
                    if (node == parent.left) {
                        setBlack(parent);
                        setRed(gparent);
                        rightRotate(gparent);
                        return;
                    }

                    if (node == parent.right) {
                        leftRotate(parent);
                        insertFixUp(parent);
                        return;
                    }
                }
            } else {
                uncle = gparent.left;

                if (uncle != null && isRed(uncle)) {
                    setBlack(parent);
                    setBlack(uncle);
                    setRed(gparent);
                    insertFixUp(gparent);
                    return;
                }

                if (uncle == null || isBlack(uncle)) {
                    if (node == parent.right) {
                        setBlack(parent);
                        setRed(gparent);
                        leftRotate(gparent);
                        return;
                    }

                    if (node == parent.left) {
                        rightRotate(parent);
                        insertFixUp(parent);
                        return;
                    }
                }
            }
        }
    }

    static class RBNode<K extends Comparable<K>, V> {
        private RBNode parent;
        private RBNode left;
        private RBNode right;
        private boolean color;
        private K key;
        private V value;

        /**
         * 获取当前节点的父节点
         */

        public RBNode() {
        }

        public RBNode(RBNode parent, RBNode left, RBNode right, boolean color, K key, V value) {
            this.parent = parent;
            this.left = left;
            this.right = right;
            this.color = color;
            this.key = key;
            this.value = value;
        }

        public RBNode getParent() {
            return parent;
        }

        public void setParent(RBNode parent) {
            this.parent = parent;
        }

        public RBNode getLeft() {
            return left;
        }

        public void setLeft(RBNode left) {
            this.left = left;
        }

        public RBNode getRight() {
            return right;
        }

        public void setRight(RBNode right) {
            this.right = right;
        }

        public boolean isColor() {
            return color;
        }

        public void setColor(boolean color) {
            this.color = color;
        }

        public K getKey() {
            return key;
        }

        public void setKey(K key) {
            this.key = key;
        }

        public V getValue() {
            return value;
        }

        public void setValue(V value) {
            this.value = value;
        }
    }

}

3.2、输出类

package p0930;

public class TreeOperation {
    /*
           树的结构示例:
              1
            /   \
          2       3
         / \     / \
        4   5   6   7
    */

    // 用于获得树的层数
    public static int getTreeDepth(RBTree.RBNode root) {
        return root == null ? 0 : (1 + Math.max(getTreeDepth(root.getLeft()), getTreeDepth(root.getRight())));
    }


    private static void writeArray(RBTree.RBNode currNode, int rowIndex, int columnIndex, String[][] res, int treeDepth) {
        // 保证输入的树不为空
        if (currNode == null) return;
        // 先将当前节点保存到二维数组中
        res[rowIndex][columnIndex] = String.valueOf(currNode.getKey() + "-" + (currNode.isColor() ? "R" : "B") + "");

        // 计算当前位于树的第几层
        int currLevel = ((rowIndex + 1) / 2);
        // 若到了最后一层,则返回
        if (currLevel == treeDepth) return;
        // 计算当前行到下一行,每个元素之间的间隔(下一行的列索引与当前元素的列索引之间的间隔)
        int gap = treeDepth - currLevel - 1;

        // 对左儿子进行判断,若有左儿子,则记录相应的"/"与左儿子的值
        if (currNode.getLeft() != null) {
            res[rowIndex + 1][columnIndex - gap] = "/";
            writeArray(currNode.getLeft(), rowIndex + 2, columnIndex - gap * 2, res, treeDepth);
        }

        // 对右儿子进行判断,若有右儿子,则记录相应的"\"与右儿子的值
        if (currNode.getRight() != null) {
            res[rowIndex + 1][columnIndex + gap] = "\\";
            writeArray(currNode.getRight(), rowIndex + 2, columnIndex + gap * 2, res, treeDepth);
        }
    }


    public static void show(RBTree.RBNode root) {
        if (root == null) System.out.println("EMPTY!");
        // 得到树的深度
        int treeDepth = getTreeDepth(root);

        // 最后一行的宽度为2的(n - 1)次方乘3,再加1
        // 作为整个二维数组的宽度
        int arrayHeight = treeDepth * 2 - 1;
        int arrayWidth = (2 << (treeDepth - 2)) * 3 + 1;
        // 用一个字符串数组来存储每个位置应显示的元素
        String[][] res = new String[arrayHeight][arrayWidth];
        // 对数组进行初始化,默认为一个空格
        for (int i = 0; i < arrayHeight; i++) {
            for (int j = 0; j < arrayWidth; j++) {
                res[i][j] = " ";
            }
        }

        // 从根节点开始,递归处理整个树
        // res[0][(arrayWidth + 1)/ 2] = (char)(root.val + '0');
        writeArray(root, 0, arrayWidth / 2, res, treeDepth);

        // 此时,已经将所有需要显示的元素储存到了二维数组中,将其拼接并打印即可
        for (String[] line : res) {
            StringBuilder sb = new StringBuilder();
            for (int i = 0; i < line.length; i++) {
                sb.append(line[i]);
                if (line[i].length() > 1 && i <= line.length - 1) {
                    i += line[i].length() > 4 ? 2 : line[i].length() - 1;
                }
            }
            System.out.println(sb.toString());
        }
    }
}

3.3、测试类

package p0930;

import java.util.Scanner;

/**
 * @author: jdq8576
 * @date: 2021-09-30 19:53
 **/
public class RBTreeTest {
    public static void main(String[] args) {
        Scanner scanner = new Scanner(System.in);
        RBTree<String, Object> rbt = new RBTree<>();

        while (true) {
            System.out.println("请输入key: ");
            String key = scanner.next();
            System.out.println();
            rbt.insert(key, null);
            TreeOperation.show(rbt.getRoot());
        }
    }
}