将BST变平衡(手撕AVL树)

力扣1382
什么是平衡二叉树(AVL)——知乎
平衡二叉树(博客园)

强调，手撕AVL并不是最优解，只是通解，时间复杂度是nlog(n)。

利用二叉搜索树的性质，中序遍历输出，然后以中间为root，递归构造树，效率更高，算是本题的最优解。

AVL的关键点不仅仅是LL/RR/LR/RL四种旋转方式，还需要注意高度的更新：1. 旋转后更新一次 2.insert递归回溯的时候更新一次，同时回溯的时候判断需要进行哪种旋转

提交时发现，必须初始化叶子节点高度为1，原因是在 height() 函数里定义空节点高度为 0，此时叶子节点必须初始化为 1；若不初始化为1,可以定义空节点高度为 -1

public AVLNode(int val){
       this.val = val;
       this.height = 1;//一定要初始化为1
 }

定义高度函数如下

private int height(AVLNode root){
      if (root == null) return 0;
      //return Math.max(height(root.left),height(root.right)) + 1;//这样写时间超时
      return root.height;
}

最后提交

//给你一棵二叉搜索树，请你返回一棵 平衡后 的二叉搜索树，新生成的树应该与原来的树有着相同的节点值。 
//
// 如果一棵二叉搜索树中，每个节点的两棵子树高度差不超过 1 ，我们就称这棵二叉搜索树是 平衡的 。 
//
// 如果有多种构造方法，请你返回任意一种。 
//
// 
//
// 示例： 
//
// 
//
// 输入：root = [1,null,2,null,3,null,4,null,null]
//输出：[2,1,3,null,null,null,4]
//解释：这不是唯一的正确答案，[3,1,4,null,2,null,null] 也是一个可行的构造方案。
// 
//
// 
//
// 提示： 
//
// 
// 树节点的数目在 1 到 10^4 之间。 
// 树节点的值互不相同，且在 1 到 10^5 之间。 
// 
// Related Topics 二叉搜索树 
// 👍 34 👎 0


//leetcode submit region begin(Prohibit modification and deletion)


import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
class Solution {
    private AVLNode head;
    private TreeNode CAST;
    class AVLNode{
        int val,height;
        AVLNode left,right;
        public AVLNode(int val){
            this.val = val;
            this.height = 1;//一定要初始化
        }
    }
    private int height(AVLNode root){
        if (root == null) return 0;
        return root.height;
    }
    //LL右旋
    private AVLNode RotateLL(AVLNode root){
        AVLNode node = root.left;
        AVLNode right = node.right;
        root.left = right;
        node.right = root;
        //更新高度，先更新右旋的root高度，在更新成为根节点的node的高度
        root.height = Math.max(height(root.left), height(root.right)) + 1;
        node.height = Math.max(height(node.left), height(node.right)) + 1;
        return node;
    }
    //RR左旋
    private AVLNode RotateRR(AVLNode root){
        AVLNode node = root.right;
        AVLNode left = node.left;
        root.right = left;
        node.left = root;
        root.height = Math.max(height(root.left), height(root.right)) + 1;
        node.height = Math.max(height(node.left), height(node.right)) + 1;
        return node;
    }
    //先左旋，再右旋
    private AVLNode RotateLR(AVLNode root){
        root.left = RotateRR(root.left);
        return RotateLL(root);
    }
    //先右旋，再左旋
    private AVLNode RotateRL(AVLNode root){
        root.right = RotateLL(root.right);
        return RotateRR(root);
    }
    private AVLNode insert(AVLNode root,int val){
        if (root == null) return new AVLNode(val);
        if (val < root.val) root.left = insert(root.left, val);
        if (val > root.val) root.right = insert(root.right, val);
        if (val == root.val) return root;
        root.height = Math.max(height(root.left), height(root.right)) + 1;//回溯更新每个节点的高度
        int balance = height(root.left) - height(root.right);//平衡因子-1,0,1
        //LL情况
        if (balance > 1 && height(root.left.left) > height(root.left.right)) return RotateLL(root);
        //LR情况
        if (balance > 1 && height(root.left.left) < height(root.left.right)) return RotateLR(root);
        //RR情况
        if (balance < -1 && height(root.right.right) > height(root.right.left)) return RotateRR(root);
        //RL情况
        if (balance < -1 && height(root.right.right) < height(root.right.left)) return RotateRL(root);
        return root;
    }
    //递归 insert 到 head 中
    private void INSERT(TreeNode root){
        if (root == null) return;
        head = insert(head, root.val);
        INSERT(root.left);
        INSERT(root.right);
    }
    //递归转换 AVLNode-->TreeNode
    private TreeNode rebuild(AVLNode root){
        if (root == null) return null;
        TreeNode node = new TreeNode(root.val);
        node.left = rebuild(root.left);
        node.right = rebuild(root.right);
        return node;
    }
    public TreeNode balanceBST(TreeNode root) {
        INSERT(root);
        CAST = rebuild(head);
        return CAST;
    }
}
//leetcode submit region end(Prohibit modification and deletion)

其中，insert函数也可以写成这样

private AVLNode insert(AVLNode root,int val){
        if (root == null) return new AVLNode(val);
        //左子树
        if (val < root.val){
            root.left = insert(root.left, val);
            if (height(root.left) - height(root.right) == 2){
                if (val < root.left.val){
                    root = RotateLL(root);
                }else {
                    root = RotateLR(root);
                }
            }
        }else if (val > root.val){//右子树
            root.right = insert(root.right, val);
            if (height(root.right) - height(root.left) == 2){
                if (val > root.right.val){
                    root = RotateRR(root);
                }else {
                    root = RotateRL(root);
                }
            }
        }else{
            return root;
        }
        //更新节点高度
        root.height = Math.max(height(root.left), height(root.right)) + 1;
        return root;
    }

这道题的最佳题解

class Solution {
    public TreeNode balanceBST(TreeNode root) {
        List<Integer> list = new ArrayList<>();
        if (root == null) return null;
        inOrder(list,root);
        if (list.size() <= 3) return root;//只剩三个以内节点，直接返回中间结点
        return helper(list,0, list.size() - 1);
    }
    private void inOrder(List<Integer> list,TreeNode root){
        if (root != null){
            inOrder(list,root.left);
            list.add(root.val);
            inOrder(list,root.right);
        }
    }
    private TreeNode helper(List<Integer> list,int l,int r){
        if (l > r) return null;
        int mid = l + (r - l) / 2;
        TreeNode root = new TreeNode(list.get(mid));
        root.left = helper(list,l, mid - 1);
        root.right = helper(list,mid + 1, r);
        return root;
    }
}

关于删除节点操作可参考《算法（第四版）》P261

   private AVLNode findMin(AVLNode root){
        if (root.left == null) return root;
        return findMin(root.left);
    }
    //删除最小节点
    private AVLNode removeMin(AVLNode root){
    //递归到左叶子节点为空的情况，返回右叶子节点
        if (root.left == null) return root.right;
        root.left = removeMin(root.left);
        root.height = Math.max(height(root.left), height(root.right)) + 1;
        return root;
    }
   private AVLNode remove(AVLNode root,int val){
        //左子树
        if (val < root.val){
            root.left = remove(root.left, val);
            if (height(root.left) - height(root.right) == 2){
                if (val < root.left.val){
                    root = RotateLL(root);
                }else {
                    root = RotateLR(root);
                }
            }
        }else if (val > root.val){//右子树
            root.right = remove(root.right, val);
            if (height(root.right) - height(root.left) == 2){
                if (val > root.right.val){
                    root = RotateRR(root);
                }else {
                    root = RotateRL(root);
                }
            }
        }else{
            if (root.left == null) return root.right;
            if (root.right == null) return root.left;
            AVLNode t = root;
            root = findMin(t.right);
            root.right = removeMin(t.right);
            root.left = t.left;
        }
        //更新节点高度
        root.height = Math.max(height(root.left), height(root.right)) + 1;
        return root;
    }