高阶 DS——手撕平衡二叉搜索树

AVL 树是一种自平衡二叉搜索树，普通的二叉搜索树存在一个问题，可能会变成一个单分支树，那么此时的查询效率就又变为了 O(n)，AVL 树的每个节点的左子树和右子树的高度差最多为 1，这确保了树的高度始终保持在对数级别，从而保证了高效的插入、删除和搜索操作，当插入一个元素打破这种平衡之后，AVL 树会通过旋转操作（单旋转或双旋转）来重新调整树的结构，使其恢复平衡。

1. AVL 树的实现

1.1. 定义节点

在二叉搜索树的基础上，加入一个平衡因子来描述左子树和右子树的高度差

public class AVLTree {
    static class TreeNode{
        public int val;
        //平衡因子，右子树高度 - 左子树高度
        public int bf;
        public TreeNode parent;
        public TreeNode left;
        public TreeNode right;

        public TreeNode(int val) {
            this.val = val;
        }
    }
    public TreeNode root;
}

1.2. 插入

刚开始的插入是和二叉搜索树一样的，都是先找到要插入的值应该插入的位置，之后再进行维护平衡

public boolean insert(int val){
TreeNode node = new TreeNode(val);
//root 为空
if(root == null){
    root = node;
    return true;
}
TreeNode cur = root;
TreeNode parent = null;
while(cur != null){
    if(val > cur.val){
        parent = cur;
        cur = cur.right;
    }else if(val < cur.val){
        parent = cur;
        cur = cur.left;
    }else{
        return false;
    }
}
//cur == null;
if(val > parent.val){
    parent.right = node;
}else{
    parent.left = node;
}
node.parent = parent;
cur = node;
}

当树为空时，那么第一个插入的元素就是根节点，之后再插入就通过 cur 来找到插入的目标位置，插入之后就需要进行平衡因子的判断，由于平衡因子是等于右子树高度 - 左子树的高度，所以当插入的位置是右子树的话，平衡因子就需要 + 1，插入的位置是左子树的话，平衡因子就 - 1，但是此时修改的只是当前位置的平衡因子，还需要不断地循环，向上调整平衡因子，如果发现有平衡因子的绝对值大于 2，就需要通过旋转来降低树的高度

1.3. 调节平衡因子

当插入元素之后，平衡因子也是要更新的，如果插入的是左树，那么平衡因子就减一，右树的话平衡因子加一

1. 如果说当前的平衡因子已经是 0 的话，无论插入左树还是右树都不会打破平衡，就不用往上更新了，直接退出循环即可
2. 如果是 1 或者 - 1的话，就代表当前的二叉树平衡了，但是上面的还不确定，所以还需要向上进行判断
3. 如果发现平衡因子是 2 或者 - 2，也就是当前二叉树不平衡了，就需要进行旋转来维持平衡，这里分为四种情况

//修改平衡因子
while (parent != null) {
    if (cur == parent.right) {
        parent.bf++;
    } else {
        parent.bf--;
    }

    if (parent.bf == 0) {
        //此时已经平衡了
        break;
    } else if (parent.bf == 1 || parent.bf == -1) { //虽然当前的二叉树平衡了，但是上面的二叉树不一定平衡
        //向上修改平衡因子
        cur = parent;
        parent = cur.parent;
    } else {
        if (parent.bf == 2) { //右树高，降低右树高度
            if (cur.bf == 1) {
                //左旋
                rotateLeft(parent);
            } else if (cur.bf == -1) {
                //右左双旋
                rotateRL(parent);
            }
        } else if (parent.bf == -2) { //左树高，降低左树高度
            if (cur.bf == 1) {
                //左右双旋
                rotateLR(parent);
            } else if (cur.bf == -1) {
                //右旋
                rotateRight(parent);
            }
        }
        //执行到这里就平衡了
        break;
    }
}

1.4. 右旋

如果说 60 是根节点的话，右旋之后需要更新 root

还可能有另外一种情况，如果 60 这个结点上面还有结点的话，又可以分为 60 是左子树和右子树两种情况

当完成右旋的话，还需要更新一下平衡因子

最后，还有一个特殊情况，如果说 subL.right 是 null 的话，还需要进行一些修改

private void rotateRight(TreeNode parent) {
        TreeNode subL = parent.left;
        TreeNode subLR = subL.right;
        parent.left = subLR;
        subL.right = parent;
        if (subLR != null) {
            subLR.parent = parent;
        }
        //记录一下parent的parent
        TreeNode pParent = parent.parent;
        parent.parent = subL;

        //如果是parent根节点，更新根节点
        if (parent == root) {
            root = subL;
            root.parent = null;
        } else {
            //判断parent结点是左子树还是右子树
            if (pParent.left == parent) {
                pParent.left = subL;
            } else if (pParent.right == parent) {
                pParent.right = subL;
            }
            subL.parent = pParent;
        }
        //修改平衡因子
        subL.bf = 0;
        parent.bf = 0;
    }

1.5. 左旋

左旋就是和刚刚的右旋相反的情况，特殊情况是和右旋类似的，也就是 parent 可能是 root 结点，也可能是左子树或者右子树，p.right.left 可能为 null 的情况

private void rotateLeft(TreeNode parent) {
        TreeNode subR = parent.right;
        TreeNode subRL = subR.left;
        parent.right = subRL;
        subR.left = parent;
        //subR.left为null的情况
        if (subRL != null) {
            subRL.parent = parent;
        }
        TreeNode pParent = parent.parent;
        parent.parent = subR;
        //parent是根节点的情况
        if (parent == root) {
            root = subR;
            root.parent = null;
        } else {
            //parent是左子树的情况
            if (pParent.left == parent) {
                pParent.left = subR;
            } else if (pParent.right == parent) {
                pParent.right = subR;
            }
            subR.parent = pParent;
        }
        //修改平衡因子
        subR.bf = 0;
        parent.bf = 0;
    }

1.6. 左右双旋

如果说是上面这种情况，那么右旋的话还是达不到平衡，左旋也达不到

可以先对 30 节点进行左旋，然后再对 60 节点进行右旋，就可以达到平衡的状态了

还有另外一种情况：

这两种情况的最终旋转完成后的平衡因子都是不一样的，根据 subLR 开始的平衡因子是 -1 还是 1 来判断最后的平衡因子怎么更新

private void rotateLR(TreeNode parent) {
TreeNode subL = parent.left;
TreeNode subLR = subL.right;
int bf = subLR.bf;
rotateLeft(parent.left);
rotateRight(parent);
if (bf == 1) {
    subLR.bf = 0;
    subL.bf = -1;
    parent.bf = 0;
} else if (bf == -1) {
    subLR.bf = 0;
    subL.bf = 0;
    parent.bf = 1;
}
}

1.7. 右左双旋

同理，仅仅通过左旋或者右旋是不能达到平衡的，需要先进行一次右旋，再进行一次左旋

根据插入位置的不同，也是有两种不同的平衡因子更新情况

private void rotateRL(TreeNode parent) {
TreeNode subR = parent.right;
TreeNode subRL = subR.left;
int bf = subRL.bf;
rotateRight(parent.right);
rotateLeft(parent);
if(bf == 1){
    subRL.bf = 0;
    subR.bf = 0;
    parent.bf = -1;
}else if(bf == -1){
    subRL.bf = 0;
    subR.bf = 1;
    parent.bf = 0;
}
}

2. 验证 AVL 树的正确性

验证时肯定要获取每一棵子树的高度，然后比较左子树和右子树的高度差，然后和计算出来的平衡因子对比，判断平衡因子有没有计算错误

private int height(TreeNode root){
if(root == null){
    return 0;
}
int leftH = height(root.left);
int rightH =height(root.right);
return Math.max(leftH,rightH) + 1;
}
public boolean isBalanced(TreeNode root){
    if(root == null){
        return true;
    }
    int leftH = height(root.left);
    int rightH = height(root.right);
    if(rightH - leftH != root.bf){
        System.out.println("该节点" + root.val + "平衡因子异常");
        return false;
    }
    return Math.abs(rightH - leftH) <= 1
    && isBalanced(root.left)
    && isBalanced(root.right);
}