Java实现AVL（平衡二叉树）

清疚 2021-09-23 07:24 339阅读 0赞

class AVLNode<T extends Comparable<T>>{
        private T data;
        private AVLNode<T> left;
        private AVLNode<T> right;
        private int height; // 记录节点当前的高度值
    
        public AVLNode(T data, AVLNode<T> left, AVLNode<T> right, int height) {
            this.data = data;
            this.left = left;
            this.right = right;
            this.height = height;
        }
    
        public T getData() {
            return data;
        }
    
        public void setData(T data) {
            this.data = data;
        }
    
        public AVLNode<T> getLeft() {
            return left;
        }
    
        public void setLeft(AVLNode<T> left) {
            this.left = left;
        }
    
        public AVLNode<T> getRight() {
            return right;
        }
    
        public void setRight(AVLNode<T> right) {
            this.right = right;
        }
    
        public int getHeight() {
            return height;
        }
    
        public void setHeight(int height) {
            this.height = height;
        }
    }
    
    /**
     * 描述:
     * @Date 2019/7/3
     */
    public class AVL<T extends Comparable<T>> {
        private AVLNode<T> root;
    
        private AVL(){
            this.root = null;
        }
    
        /**
         * 以参数node为根节点进行左旋操作，把旋转后的树的根节点返回
         * @param node
         * @return
         */
        private AVLNode<T> leftRotate(AVLNode<T> node){
            AVLNode<T> child = node.getRight();
            node.setRight(child.getLeft());
            child.setLeft(node);
    
            //更新结点的高度
            node.setHeight(maxHeight(node.getLeft(),node.getRight())+1);
            child.setHeight(maxHeight(child.getLeft(),child.getRight())+1);
    
            return child;
        }
    /*----------------------------------------------------------------------------------------------------------------------------------------*/
    
        /**
         * 以参数node为根节点进行右旋操作，把旋转后的树的根节点返回
         * @param node
         * @return
         */
        private AVLNode<T> rightRotate(AVLNode<T> node){
            AVLNode<T> child = node.getLeft();
            node.setLeft(child.getRight());
            child.setRight(node);
    
            //更新结点的高度
            node.setHeight(maxHeight(node.getLeft(),node.getRight())+1);
            child.setHeight(maxHeight(child.getLeft(),child.getRight())+1);
    
            return child;
        }
    
    /*----------------------------------------------------------------------------------------------------------------------------------------*/
    
        /**
         * 节点的高度
         **/
        private int maxHeight(AVLNode<T> left, AVLNode<T> right) {
            int l = height(left);
            int r = height(right);
            return l > r ? 1 : r;
        }
    
        private int height(AVLNode<T> root) {
            return root == null ? 0: root.getHeight();
        }
    
    /*----------------------------------------------------------------------------------------------------------------------------------------*/
    
        /**
         * 以参数node为根节点进行左平衡操作，把旋转后的树的根节点返回
         * 左-右旋转
         * @param node
         * @return
         */
        private AVLNode<T> leftBalance(AVLNode<T> node){
            node.setLeft(leftRotate(node.getLeft()));   //先以根节点的左孩子为根节点左旋，更新一下根节点的左孩子
            return rightRotate(node); //以更新后的根节点进行右旋
        }
    
    /*----------------------------------------------------------------------------------------------------------------------------------------*/
        /**
         * 以参数node为根节点进行右平衡操作，把旋转后的树的根节点返回
         * 右-左旋
         * @param node
         * @return
         */
        private AVLNode<T> rightBalance(AVLNode<T> node){
            node.setRight(rightRotate(node.getRight()));  //先以根节点的右孩子为根节点右旋，更新一下根节点的右孩子
            return leftRotate(node);
        }
    
    /*----------------------------------------------------------------------------------------------------------------------------------------*/
        /**
         * 以参数node为起始节点，搜索一个合适的位置添加data，再把子树的根节点返回
         * 递归实现AVL树的插入操作
         * @param data
         * @return AVLNode<T>
         */
        public void insert(T data){
            this.root = insert(this.root,data);
        }
    
        private AVLNode<T> insert(AVLNode<T> root,T data) {
    
            if (root == null) {
                return new AVLNode<>(data, null, null, 1);
            }
    
            if (root.getData().compareTo(data) > 0) {
                root.setLeft(insert(root.getLeft(), data));
                //判断root节点是否失衡
                if (height(root.getLeft()) - height(root.getRight()) > 1) { //左 > 右 失衡
                    //左孩子的左子树太高，右旋
                    if (height(root.getLeft().getLeft()) >= height(root.getLeft().getRight())) {
                        root = rightRotate(root);
                    } else {
                        //左孩子的右子树太高，左平衡
                        root = leftBalance(root);
                    }
                }
            } else if (root.getData().compareTo(data) < 0) {
                root.setRight(insert(root.getRight(), data));
                //判断root节点是否失衡
                if (height(root.getRight()) - height(root.getLeft()) > 1) { //右 > 左 失衡
                       //右孩子的右子树太高,右旋
                    if (height(root.getRight().getRight()) >= height(root.getRight().getLeft())) {
                        root = leftRotate(root);
                    } else {
                        //右孩子的左子树太高,右平衡
                        root = rightBalance(root);
                    }
                }
            }
            //递归回溯的过程中，更新节点的高度值
            root.setHeight(maxHeight(root.getLeft(), root.getRight()) + 1);
             return root;
            }
    
    /*----------------------------------------------------------------------------------------------------------------------------------------*/
        /**
         * AVL树的递归删除操作
         * @param data
         * */
            public void remove(T data){
                this.root = remove(this.root,data);
            }
    
           private AVLNode<T> remove(AVLNode<T> root, T data) {
               if (root == null) {
                   return null;
               }
               if (root.getData().compareTo(data) > 0) {  //在左边找
                   root.setLeft(remove(root.getLeft(), data));
    
                   //判断root节点是否失衡
                   if (height(root.getLeft()) - height(root.getRight()) > 1) { //左 > 右 失衡
                       //左孩子的左子树太高，右旋
                       if (height(root.getLeft().getLeft()) >= height(root.getLeft().getRight())) {
                           root = rightRotate(root);
                       } else {
                           //左孩子的右子树太高，左平衡
                           root = leftBalance(root);
                       }
                   }
               } else if (root.getData().compareTo(data) < 0) {
                   root.setRight(insert(root.getRight(), data));
                   //判断root节点是否失衡
                   if (height(root.getRight()) - height(root.getLeft()) > 1) { //右 > 左 失衡
                       //右孩子的右子树太高,右旋
                       if (height(root.getRight().getRight()) >= height(root.getRight().getLeft())) {
                           root = leftRotate(root);
                       } else {
                           //右孩子的左子树太高,右平衡
                           root = rightBalance(root);
                       }
                   }
               } else if (root.getData().compareTo(data) < 0) {  //在右边找
                   root.setRight(remove(root.getRight(), data));
               } else {
                   if (root.getLeft() != null && root.getRight() != null) {
                       if(height(root.getLeft()) >= height(root.getRight())){  //左大于右 高度
                           //用前驱替换
                           AVLNode<T> pre = root.getLeft();
                           while (pre.getRight() != null) {
                               pre = pre.getRight();
                           }
                           root.setData(pre.getData());
                           root.setLeft(remove(root.getLeft(), pre.getData())); //删除前驱
                       }else {
                           //用后继替换
                           AVLNode<T> post = root.getRight();
                           if (root.getLeft() != null) {
                               post = post.getRight();
                           }
                           root.setData(post.getData());
                           root.setRight(remove(root.getRight(),post.getData())); //删除后继
                       }
                   }else{
                       if(root.getLeft() != null){
                           return root.getLeft();
                       }else if(root.getRight() != null){
                           return root.getRight();
                       }else {
                           return null;
                       }
                   }
               }
    
               return root;
           }
    
    
    
           public static void main(String[] args) {
            AVL<Integer> avl = new AVL<>();
            for(int i = 1; i < 11; i++){
                avl.insert(i);
            }
         }
    
    }