数据结构(十)：AVL树主体封装

封装AVL树主体

继承常规二叉搜索树

由于AVL树也是二叉搜索树，所以很多方法是可以继承常规的二叉搜索树的，我们在前面章节已经封装过二叉搜索树，所以这里我们直接继承即可。

class AVLTree<T> extends BSTree<T> {
  
}

由于要使用常规二叉搜索树的插入和删除方法，对于AVL节点也应该继承常规二叉搜索树的节点：

// AVl树节点
class AVLTreeNode<T> extends TreeNode<T> {
  left: AVLTreeNode<T> | null = null;
  right: AVLTreeNode<T> | null = null;
  parent: AVLTreeNode<T> | null = null;
  // 高度
  get height(): number {
    ...
  }
  // 平衡因子 = 左子树高度 - 右子树高度
  get balanceFactor(): number {
    ...
  }
  // 是否是平衡节点
  get isBalanced(): boolean {
    ...
  }
  // 获取子节点中高度较高的节点 该节点就是要旋转的轴心节点
  get higherChild(): AVLTreeNode<T> | null {
    ...
  }
  // 左右旋转封装 传入旋转方向
  rotateByDir(direction: 'left' | 'right') {
    ...
  }
}

特殊节点失衡情况分析

上一章封装了AVL树节点以及处理了左左情况和右右情况，其实除了左左情况和右右情况，还有左右情况和右左情况，如下图：

先来分析左右情况：遇到左右情况时我们发现如果还按照左左或者右右的处理方式是不行的，但是如果我们按照将3当成root、把4当成pivot先进行一次左旋转就会的到如下图：

可以看见变成了左左的情况，此时再进行一次右旋转就能得到平衡的树结构了：

那么对于右左的情况也是类似的，我们先进行一次右旋转，再进行一次左旋转就可以了。
到此就是节点失衡的所有会出现的情况，维基百科上的一幅图可以非常直观的总结了这四种情况：

rebalance封装

假如现在我们已经找到了失衡节点，那么我们要区分上述四种不同的情况，去执行不同的旋转，我们封装一个rebalance方法来实现这个目的：

class AVLTree<T> extends BSTree<T> {
  // 重新平衡非平衡节点
  rebalance(unbalancedNode: AVLTreeNode<T>) {
    // 获取pivot节点
    const pivot = unbalancedNode.higherChild!;
    // 获取导致不平衡的节点
    const theCauseNode = pivot.higherChild!;
    // 如果旋转后pivot没有父节点 则将其设置为根节点
    let newRoot: AVLTreeNode<T> | null = null;
    // 区分四种情况 LL LR RL RR
    // 如果pivot是左子节点 一定是LL或者LR
    if (pivot.isLeftChild) {
      if (theCauseNode.isLeftChild) {
        // LL 右旋转
        newRoot = unbalancedNode.rotateByDir('right');
      } else {
        // LR 先左旋转 以pivot为root theCauseNode为轴心 
        pivot.rotateByDir('left');
        // 再右旋转 以unbalancedNode为root pivot为轴心 
        newRoot = unbalancedNode.rotateByDir('right');
      }
    } else {
      // 如果pivot是右子节点 一定是RL或者RR
      if (theCauseNode.isLeftChild) {
        // RL 先右旋转 以pivot为root theCauseNode为轴心 
        pivot.rotateByDir('right');
        // 再左旋转 以unbalancedNode为root pivot为轴心 
        newRoot = unbalancedNode.rotateByDir('left');
      } else {
        // RR 左旋转
        newRoot = unbalancedNode.rotateByDir('left');
      }
    }
    // 如果newRoot存在 则将其设置为根节点
    if (newRoot) this.root = newRoot;
  }
}

可以看到代码实现非常简单，主要是要理清楚思路

插入操作

插入操作可能会导致某个节点不平衡，所以要在以前二叉搜索树的插入操作基础上添加上判断是否有节点不平衡，找出不平衡节点并调用rebalance方法的逻辑。
这里注意要复用BSTree的insert方法时，由于其内部使用的是TreeNode类创建子节点，所以我们要适当改造以便AVL能使用自己的子节点类，只需要BSTree新增一个创建节点的方法，然后AVL重写它即可：

export class BSTree<T> {
  // 创建节点
  protected createTreeNode(value: T): TreeNode<T> {
    return new TreeNode(value);
  }
  // 插入
  insert(value: T, defaultTreeNode = TreeNode) {
    const node = this.createTreeNode(value);
    ...
  }
}
class AVLTree<T> extends BSTree<T> {
  // 重写创建节点方法
  protected createTreeNode(value: T): TreeNode<T> {
    return new AVLTreeNode(value);
  }
}

那么插入节点后怎么判断有没有不平衡的节点呢？其实非常简单，只需要不断判断父节点是否平衡直到根节点，如果到根节点都是平衡的，则没有不平衡节点。对于获取当前节点有两种方式，一种是BSTree方法的insert增加返回值；另一种是新增一个方法，insert插入后调该方法，AVL中重写该方法。这两种都是可以的，我们这里使用第二种：

export class BSTree<T> {
  // 检查是否平衡
  checkBalance(node: TreeNode<T>) {}
  // 插入
  insert(value: T, defaultTreeNode = TreeNode) {
    const node = this.createTreeNode(value);
    ...
    this.checkBalance(node);
  }
}
class AVLTree<T> extends BSTree<T> {
  // 实现检查是否平衡方法
  checkBalance(node: AVLTreeNode<T>) {
    let cur = node.parent;
    // 从当前节点的父节点开始向上遍历
    while (cur) {
      // 如果不平衡 则进行平衡操作
      if (!cur.isBalanced) {
        this.rebalance(cur);
        // 插入操作只会导致一个节点不平衡 所以这里可以直接break
        break;
      }
      cur = cur.parent;
    }
  }
}

删除操作

看下图：

如果移除12和25将导致50节点不平衡，所以要平衡50节点得到右边的树。跟插入操作类似，删除节点后我们也要找不平衡节点，删除比插入时的寻找操作要复杂，大致分为两种情况：

• 当删除节点是叶子节点或者有一个子节点，那么直接从删除节点位置的父节点向上查找即可
• 当删除节点有两个子节点，要从前驱节点的父节点或者后置节点的父节点向上查找，这里还有一种特殊情况即前驱节点或者后置节点的父元素是要被删除的节点，此时要从被删除的位置开始查找。

那么按照这个思路去实现代码（有两个子节点的情况这里使用后置节点）：

export class BSTree<T> {
  // 删除
  remove(value: T) {
    // 检查节点是否平衡的开始节点
    let checkBlcStart: TreeNode<T> | null = null;
    ...
    // 判断removeNode是否是叶子节点
    if (removeNode.left === null && removeNode.right === null) {
      ...
      checkBlcStart = removeNode.parent;
    }
    // 删除节点只有一个子节点
    else if (removeNode.left === null || removeNode.right === null) {
      ...
      checkBlcStart = removeNode.parent;
      // 处理removeNode子节点的父节点
      (removeNode.left || removeNode.right)!.parent = removeNode.parent;
    }
    // 删除节点有两个子节点
    else {
      // 获取后继节点
      const successor = this.getSuccessor(removeNode);
      // 如果后置节点就是removeNode的右节点 checkBlcStart需要赋值为successor
      // 因为removeNode要被删除了 不能参与到旋转操作中
      if (successor === removeNode.right) {
        checkBlcStart = successor
      } else {
        checkBlcStart = successor.parent
      }
      if (removeNode.parent) {
        // 判断removeNode是否是父节点的左子节点
        if (removeNode.isLeftChild) {
          removeNode.parent.left = successor;
        } else {
          removeNode.parent.right = successor;
        }
      } else {
        // 如果是根节点
        this.root = successor;
      }
      // 处理后继节点的父节点
      successor.parent = removeNode.parent;
    }
    // 删除完成后 检查树平衡
    this.checkBalance(checkBlcStart, false);
  }
}
class AVLTree<T> extends BSTree<T> {
  // 实现检查是否平衡
  checkBalance(node: AVLTreeNode<T> | null, isStartFromParent = true) {
    let cur = isStartFromParent ? node!.parent : node;
    // 从当前节点的父节点开始向上遍历
    while (cur) {
      // 如果不平衡 则进行平衡操作
      if (!cur.isBalanced) {
        this.rebalance(cur);
        // 插入只会导致一个节点不平衡 直接break即可
        if (isStartFromParent) break;
        // 删除可能会导致多个节点不平衡 所以需要继续向上遍历
      }
      cur = cur.parent;
    }
  }
}

注意获取后置节点方法中的successor.parent = removeNode.parent;操作被提出来放到后面执行，因为checkBlcStart要记录后驱节点原本的parent。还有checkBalance操作因为删除节点记录的是父节点，所以要从传入节点开始遍历，所以增加一个入参来决定从什么地方开始遍历，同时为false时也代表是删除时调用此方法。
这里有一个细节是checkBalance中while循环里的break只有插入的时候可以加，删除操作的时候是不能break的，这是为什么呢？先思考插入的时候为什么可以加。因为插入前跟插入旋转后以不平衡节点为根节点的树的高度是不变的，那么对于不平衡节点的父元素来说该子树高度没有变化，自然就不会有其他不平衡的节点产生，所以不用继续向上找。而删除不一样，删除操作是会改变以不平衡节点为根的树的高度，如下图所示：

以50为根节点的树的高度由3变成2了，所以会影响50节点以上节点的不平衡判断，所以不能break。
以上就是AVL的全部封装了，细节还是比较多的，但是只要把思路理清楚，写出来也是不难的，完整的代码：参考代码，如果你有收获，记得点赞关注，后续更精彩！