午休时间，手写了红黑树，转啊转~变啊变~

ゝ一世哀愁。 2023-10-07 21:32 29阅读 0赞

package com.ossa.issavior.binarysearchtree;
    
    /**
     * 红黑树
     *
     * @param <K> K
     * @param <V> V
     */
    public class RBSTreeMap<K extends Comparable<K>, V> {
        
    
        /**
         * 原始根节点
         */
        private RBSTreeNode<K, V> firstRoot;
    
        /**
         * 节点左旋
         *
         * @param node 当前节点
         */
        private void leftRotate(RBSTreeNode<K, V> node) {
        
    
            if (node.right == null) {
        
                return;
            }
    
            if (node.parent.left == node) {
        
                node.parent.left = node.right;
            } else if (node.parent.right == node) {
        
                node.parent.right = node.right;
            } else firstRoot = node.right;
    
            node.right.parent = node.parent;
            node.parent = node.right;
            if (node.right.left != null) {
        
                node.right = node.right.left;
            }
            node.right.left = node;
        }
    
        /**
         * 节点右旋
         *
         * @param node 当前节点
         */
        private void rightRotate(RBSTreeNode<K, V> node) {
        
            if (node.left == null) {
        
                return;
            }
    
            if (node.parent.left == node) {
        
                node.parent.left = node.left;
            } else if (node.parent.right == node) {
        
                node.parent.right = node.left;
            } else firstRoot = node.left;
    
            node.left.parent = node.parent;
            node.parent = node.left;
            if (node.left.right != null) {
        
                node.left = node.left.right;
            }
            node.left.right = node;
        }
    
    
        /**
         * 根据Key搜索某一个节点
         *
         * @param key 需要搜索的值
         * @return 节点值
         */
        public V get(K key) {
        
            // 如果节点不存在，则返回空
            if (indexSearch(firstRoot, key) == null) {
        
                return null;
            }
            return indexSearch(firstRoot, key).nodeValue;
        }
    
        /**
         * 存值
         *
         * @param key   key
         * @param value value
         */
        public void put(K key, V value) {
        
    
            // 将key和value封装成node对象
            RBSTreeNode<K, V> newNode = new RBSTreeNode<K, V>(key, value);
            insertNode(firstRoot, newNode);
        }
    
        /**
         * 根据key删除node节点
         *
         * @param key key
         */
        public void delete(K key) {
        
            removeNode(key);
        }
    
        /**
         * 新增node
         *
         * @param root    root
         * @param newNode 新增的节点
         */
        void insertNode(RBSTreeNode<K, V> root, RBSTreeNode<K, V> newNode) {
        
    
            newNode.color = Color.RED;
    
            // 如果当前根节点为null，将当前根节点引用指向新的node节点
            if (root == null) {
        
                firstRoot = newNode;
                firstRoot.color = Color.BLACK;
                return;
            }
    
            // 如果当前根节点key大于新增的节点key，则向当前节点的左节点进行操作
            if (root.nodeKey.compareTo(newNode.nodeKey) > 0) {
        
                // 如果当前根节点的左节点为null，则将当前根节点的左节点引用指向新增的节点，
                // 否则进行递归，直到左节点尽头为止
                if (root.left == null) {
        
                    root.left = newNode;
                    newNode.parent = root;
                    correct(newNode);
                } else insertNode(root.left, newNode);
            } else if (root.nodeKey.compareTo(newNode.nodeKey) < 0) {
        
                // 如果当前根节点key小于新增的节点key，则向当前节点的右节点进行操作，
                // 如果右节点为空，则将当前根节点的右节点引用指向新增的节点，
                // 否则进行递归，直到右节点尽头为止
                if (root.right == null) {
        
                    root.right = newNode;
                    newNode.parent = root;
                    correct(newNode);
    
                } else insertNode(root.right, newNode);
            }
    
        }
    
        /**
         * 旋转和变色，修正红黑树
         *
         * @param newNode 节点
         */
        private void correct(RBSTreeNode<K, V> newNode) {
        
    
            // 插入的节点为红色，如果当前节点的父节点是红色，则分三种情况
            while (newNode.parent != null && newNode.parent.color == Color.RED) {
        
    
                if (newNode.parent.parent.left == newNode.parent) {
        
    
                    // 如果当前节点的叔叔节点为红色，那么则需要，将父节点变为黑色，叔叔节点也要变为黑色，
                    // 如果祖父不为根节点则将其变为红色，然后将祖父节点当作新插入的节点继续进行相关操作
                    if (newNode.parent.parent.right == null || newNode.parent.parent.right.color == Color.BLACK) {
        
                        if (newNode.parent.right == newNode) {
        
                            leftRotate(newNode.parent);
                        } else rightRotate(newNode.parent);
                    } else {
        
                        newNode.parent.color = Color.BLACK;
                        newNode.parent.parent.right.color = Color.BLACK;
                        newNode.parent.parent.color = Color.RED;
                    }
                } else {
        
    
                    if (newNode.parent.parent.left == null || newNode.parent.parent.left.color == Color.BLACK) {
        
                        if (newNode.parent.right == newNode) {
        
                            leftRotate(newNode.parent);
                        } else rightRotate(newNode.parent);
                    } else {
        
                        newNode.parent.color = Color.BLACK;
                        newNode.parent.parent.left.color = Color.BLACK;
                        newNode.parent.parent.color = Color.RED;
                    }
                }
            }
        }
    
    
        /**
         * 删除node
         *
         * @param key key
         */
        private void removeNode(K key) {
        
    
            // 根据key找到节点所在的位置
            RBSTreeNode<K, V> removeTargetNode = indexSearch(firstRoot, key);
            // 如果找不到，直接返回
            if (removeTargetNode == null) {
        
                return;
            }
    
            // 将要删除的节点的父节点
            RBSTreeNode<K, V> parent = removeTargetNode.parent;
            // 如果父节点为空，说明没有节点或者只有一个根节点，则将根节点置为空即可
            // 否则，将分为3种情况：
            //     1. 删除的节点没有左右节点，直接删除即可
            //     2. 删除的节点有左节点，没有右节点，那么需要将删除节点的父节点的左/右节点指向删除节点的左节点，
            //        删除节点的左节点的父节点需要指向删除节点的父节点
            //     3. 删除的节点有右节点，没有左节点，那么需要将删除节点的父节点的左/右节点指向删除节点的右节点，
            //        删除节点的右节点的父节点需要指向删除节点的父节点
            //     4. 删除的节点有左右节点，那么需要找到当前节点的中序前序节点或中序后继节点，替换掉当前节点，包括key和value，
            //        再删掉这个当前节点的中序前序节点或中序后继节点
            if (parent == null) {
        
                firstRoot = null;
            } else {
        
                //     1. 删除的节点没有左右节点，直接删除即可
                if (removeTargetNode.left == null && removeTargetNode.right == null) {
        
    
                    if (parent.left == removeTargetNode) {
        
                        parent.left = null;
                    } else if (parent.right == removeTargetNode) {
        
                        parent.right = null;
                    }
                }
    
                //     2. 删除的节点有左节点，没有右节点，那么需要将删除节点的父节点的左/右节点指向删除节点的左节点，
                //        删除节点的左节点的父节点需要指向删除节点的父节点
                if (removeTargetNode.left != null) {
        
                    if (parent.left == removeTargetNode) {
        
                        parent.left = removeTargetNode.left;
                        removeTargetNode.left.parent = parent;
    
                    } else if (parent.right == removeTargetNode) {
        
                        parent.right = removeTargetNode.left;
                        removeTargetNode.left.parent = parent;
                    }
    
                    if (removeTargetNode.color == Color.BLACK)
                        repair(removeTargetNode.left);
    
                }
                //     3. 删除的节点有右节点，没有左节点，那么需要将删除节点的父节点的左/右节点指向删除节点的右节点，
                //        删除节点的右节点的父节点需要指向删除节点的父节点
                if (removeTargetNode.right != null) {
        
                    if (parent.left == removeTargetNode) {
        
                        parent.left = removeTargetNode.right;
                        removeTargetNode.right.parent = parent;
    
                    } else if (parent.right == removeTargetNode) {
        
                        parent.right = removeTargetNode.right;
                        removeTargetNode.right.parent = parent;
                    }
    
                    if (removeTargetNode.color == Color.BLACK)
                        repair(removeTargetNode.right);
                }
    
                //     4. 删除的节点有左右节点，那么需要找到当前节点的中序前序节点或中序后继节点，替换掉当前节点，包括key和value，
                //        再删掉这个当前节点的中序前序节点或中序后继节点
                if (removeTargetNode.left != null && removeTargetNode.right != null) {
        
    
                    RBSTreeNode<K, V> left = removeTargetNode.left;
                    while (left.right != null) {
        
                        left = left.right;
                    }
                    removeTargetNode.nodeKey = left.nodeKey;
                    removeTargetNode.nodeValue = left.nodeValue;
                    removeNode(left.nodeKey);
                }
            }
        }
    
        /**
         * 因为删除节点而导致树的结构发生变化，需要修复树的结构
         *
         * @param targetNode 目标的节点
         */
        private void repair(RBSTreeNode<K, V> targetNode) {
        
    
            if (targetNode.color == Color.RED)
                targetNode.color = Color.BLACK;
    
            if (targetNode.color == Color.BLACK) {
        
                // 如果不是根节点
                if (targetNode != firstRoot) {
        
                    // 4种case
                }
            }
        }
    
        /**
         * 给定一个根节点，根据Key搜索某一个节点
         * 时间复杂度和二叉树的层数有关，所以时间复杂度为O(Logn)
         *
         * @param root 根节点引用
         * @param key  需要搜索的值
         * @return 节点
         */
        private RBSTreeNode<K, V> indexSearch(RBSTreeNode<K, V> root, K key) {
        
            //如果当前的根节点为null，返回空节点
            if (root == null) {
        
                return null;
            }
            // 如果当前节点key大于待搜索值，那么向左查找，反之，向右查找。
            if (root.nodeKey.compareTo(key) > 0) {
        
                return indexSearch(root.left, key);
            } else if (root.nodeKey.compareTo(key) < 0) {
        
                return indexSearch(root.right, key);
            } else return root;
        }
    
        /**
         * 内部类-node
         *
         * @param <K> K
         * @param <V> V
         */
        static class RBSTreeNode<K, V> {
        
    
            /**
             * 树的颜色：true-黑；false-红
             */
            private Color color;
    
            /**
             * 键值
             */
            private K nodeKey;
    
            /**
             * 值
             */
            private V nodeValue;
    
            /**
             * 左节点引用
             */
            private RBSTreeNode<K, V> left;
    
            /**
             * 右节点引用
             */
            private RBSTreeNode<K, V> right;
    
            /**
             * 父节点引用
             */
            private RBSTreeNode<K, V> parent;
    
            public RBSTreeNode(K nodeKey, V nodeValue) {
        
                this.nodeKey = nodeKey;
                this.nodeValue = nodeValue;
    
            }
    
            @Override
            public String toString() {
        
                return "RBSTreeNode{" +
                        "nodeKey=" + nodeKey +
                        '}';
            }
    
        }
    
        /**
         * 红黑树的颜色
         */
        enum Color {
        
            BLACK,
            RED;
        }
    }