TypeScript数据结构与算法：AVL树

上一篇《TypeScript 数据结构与算法：二叉搜索树》实现了 Typescript 中二叉搜索树的数据结构与算法，本篇继续进一步实现自平衡的AVL树。

上一篇中实现的二叉搜索树是最常用的二叉树，支持快速的插入、删除、查找等操作，操作的时间复杂度跟树的高度成正比，所以在理想情况下（比如完全二叉树），时间复杂度是 O(logn)。

但是如果频繁的操作二叉搜索树之后，有可能会出现树的一条边特别深的情况，如下图所示：

此时的树高度就可能比 log₂n 大得多，时间复杂度也会随之变大。在极端情况下，二叉树就退化成了链表，时间复杂度退化为了 O(n)。为了解决这个问题，就需要二叉搜索树具有自平衡的功能，本篇实现的 AVL 树就是一种典型的自平衡树。

AVL 树是最早被发明的自平衡树，得名于它的发明者 G. M. Adelson-Velsky 和 Evgenii Landis，他们在 1962 年的论文《An algorithm for the organization of information》中公开了这一数据结构。AVL 树在添加或移除节点时会尝试保持自平衡。任意一个节点（不论深度）的左子树和右子树高度最多相差 1。添加或移除节点时，AVL 树会尽可能尝试转换为完全树。

继承二叉查找树

AVL 树也是一个二叉搜索树，所以可以直接继承上篇实现的二叉搜索树 BinarySearchTree：

import { Compare, defaultCompare, ICompareFunction } from "../util";
import BinarySearchTree from "./binary-search-tree";
import { Node } from "./models/node";


// 继承了上一节实现的二叉搜索树
export default class AVLTree<T> extends BinarySearchTree<T> {
  constructor(protected compareFn: ICompareFunction<T> = defaultCompare) {
    super(compareFn);
  }

  // ......
}

节点高度和平衡因子

节点的高度是从节点到其任意子节点的边的最大数目。下图展示了一个包含每个节点高度的树：

节点的左子树与右子树的高度差即为该节点的平衡因子（BF,Balance Factor）。

当一棵二叉搜索树为平衡二叉树时，树上所有节点的平衡因子只可能是-1（轻微右重），0（平衡） 或 1（轻微左重）。另外，由于每次插入和删除时都会进行再平衡操作，所以在 AVL 树算法中的平衡因子不会超过 -2（右重） 或 2（左重）。

// 平衡因子枚举
enum BalanceFactor {
  UNBALANCED_RIGHT = -2, // 右重
  SLIGHTLY_UNBALANCED_RIGHT = -1, // 轻微右重
  BALANCED = 0, // 平衡
  SLIGHTLY_UNBALANCED_LEFT = 1, // 轻微左重
  UNBALANCED_LEFT = 2 // 右重
}

/**
  * @description: 获取节点高度
  */
private getNodeHeight(node: Node<T>): number {
  // 基线条件
  if (node == null) {
    return -1;
  }
  // 递归计算
  return (
    Math.max(this.getNodeHeight(node.left), this.getNodeHeight(node.right)) +
    1
  );
}

/**
  * @description: 获取节点的平衡因子
  * @param {Node} node
  */
private getBalanceFactor(node: Node<T>): BalanceFactor {
  // 左子树重 减去 右子树重
  const heightDifference =
    this.getNodeHeight(node.left) - this.getNodeHeight(node.right);
  // 再返回对应的枚举值
  switch (heightDifference) {
    case -2:
      return BalanceFactor.UNBALANCED_RIGHT;
    case -1:
      return BalanceFactor.SLIGHTLY_UNBALANCED_RIGHT;
    case 1:
      return BalanceFactor.SLIGHTLY_UNBALANCED_LEFT;
    case 2:
      return BalanceFactor.UNBALANCED_LEFT;
    default:
      return BalanceFactor.BALANCED;
  }
}

平衡旋转

AVL 树实现的原理就是，在对树添加或移除节点后，计算各节点的高度并验证树是否需要进行平衡。而 AVL 实现平衡的核心方式就是旋转。

旋转时可以执行单旋转或双旋转两种平衡操作，分别对应四种场景：

• 左-左（LL）：右旋
• 右-右（RR）：左旋
• 左-右（LR）：先右旋后左旋（先 LL 旋转，再 RR 旋转）
• 右-左（RL）：先左旋后右旋（先 RR 旋转，再 LL 旋转）

先来看下单旋转的实现方式，下图分别是对 A 节点进行右旋和对 B 节点进行左旋的动图：

下图为全部四种情况的旋转对比：

/**
  * 左左情况: 向右单旋转
  * 左侧子节点的高度大于右侧子节点的高度时，并且左侧子节点也是平衡或左侧较重的，为左左情况
  *
  *       a                           b
  *      / \                         / \
  *     b   c -> rotationLL(a) ->   d   a
  *    / \                             / \
  *   d   e                           e   c
  *
  * @param root Node<T> 旋转前的root节点
  * @returns pivot Node<T> 返回旋转后的root节点(也就是旋转前的pivot)
  */
private rotationLL(root: Node<T>): Node<T> {
  // 先把pivot拿出来
  const pivot = root.left;
  // root 左侧指向 pivot 的右子
  root.left = pivot.right;
  // pivot 右侧指向 root
  pivot.right = root;
  // 返回 pivot
  return pivot;
}

/**
  * 右右情况: 向左单旋转
  * 右侧子节点的高度大于左侧子节点的高度，并且右侧子节点也是平衡或右侧较重的，为右右情况
  *     a                              c
  *    / \                            / \
  *   b   c   -> rotationRR(a) ->    a   e
  *      / \                        / \
  *     d   e                      b   d
  *
  * @param root Node<T> 旋转前的root节点
  * @returns pivot Node<T> 返回旋转后的root节点(也就是旋转前的pivot)
  */
private rotationRR(root: Node<T>): Node<T> {
  // 先把pivot拿出来
  const pivot = root.right;
  // root 右侧指向 pivot 的左子
  root.right = pivot.left;
  // pivot 左侧指向 root
  pivot.left = root;
  // 返回 pivot
  return pivot;
}

/**
  * 左右情况: 先左转子节点后右转
  * 左侧子节点的高度大于右侧子节点的高度，并且左侧子节点右侧较重
   *
   *       a                           a                            e
   *      / \                         / \                         /   \
   *     b   c -> rotationRR(b) ->   e   c -> rotationLL(a) ->   b     a
   *    / \                         / \                         / \   / \
   *   d   e                       b   g                       d   f g   c
   *      / \                     / \
   *     f   g                   d   f
   *
  * @param node Node<T>
  */
private rotationLR(node: Node<T>): Node<T> {
  // 先把节点左子左转
  node.left = this.rotationRR(node.left);
  // 再把节点本身右转
  return this.rotationLL(node);
}

/**
  * 右左情况: 先右转子节点后左转
  * 右侧子节点的高度大于左侧子节点的高度，并且右侧子节点左侧较重
  *
  *       a                           a                            d
  *      / \                         / \                         /   \
  *     b   c -> rotationLL(c) ->   b   d -> rotationRR(a) ->   a     c
  *        / \                         / \                     / \   / \
  *       d   e                       f   c                   b   f g   e
  *      / \                             / \
  *     f   g                           g   e
  *
  * @param node Node<T>
  */
private rotationRL(node: Node<T>): Node<T> {
  // 先把节点右子右转
  node.right = this.rotationLL(node.right);
  // 再把节点本身左转
  return this.rotationRR(node);
}

保持平衡

前面实现的是进行旋转的具体步骤，下面将各个旋转对应的各个情况汇总，写成一个单独的方法，方便复用：

/**
  * @description: 对节点为根的树进行两层平衡
  * @param {Node} node
  * @return {Node} 返回平衡后的以节点为根的树
  */
keepBalance(node: Node<T>): Node<T> {
  // 先校验tree是否是平衡的，只有“左重”和“右重”时才需要重新再平衡，
  // “轻微左重”、“轻微右重”、“平衡”的三种状态不需要再平衡
  if (node == null) {
    return node;
  }
  // 校验树是否平衡
  const balanceState = this.getBalanceFactor(node);

  if (balanceState === BalanceFactor.UNBALANCED_LEFT) {
    // 左左情况
    if (
      this.getBalanceFactor(node.left) === BalanceFactor.BALANCED ||
      this.getBalanceFactor(node.left) ===
        BalanceFactor.SLIGHTLY_UNBALANCED_LEFT
    ) {
      return this.rotationLL(node);
    }
    // 左右情况
    else if (
      this.getBalanceFactor(node.left) ===
      BalanceFactor.SLIGHTLY_UNBALANCED_RIGHT
    ) {
      return this.rotationLR(node);
    }
  } else if (balanceState === BalanceFactor.UNBALANCED_RIGHT) {
    // 右右情况
    if (
      this.getBalanceFactor(node.right) === BalanceFactor.BALANCED ||
      this.getBalanceFactor(node.right) ===
        BalanceFactor.SLIGHTLY_UNBALANCED_RIGHT
    ) {
      return this.rotationRR(node);
    }
    // 右左情况
    else if (
      this.getBalanceFactor(node.right) ===
      BalanceFactor.SLIGHTLY_UNBALANCED_LEFT
    ) {
      return this.rotationRL(node);
    }
  }

  // “轻微左重”、“轻微右重”、“平衡”的三种状态不需要再平衡，直接返回
  return node;
}

重写插入和删除

最后将二叉搜索树中的插入和删除``覆写，其实就是在最后执行了 keepBalance 操作。

需要注意的是，由于 insertNode 和 removeNode 都是一个递归操作，所以也就是说，每一个与添加和删除有关的节点都会执行一次 keepBalance，最终实现了整个树的平衡。

/**
  * @description: 插入节点的递归方法，递归插入完后，需要校验树是否仍然平衡，若不平衡则需要旋转
  * @param {Node} node 要插入到的节点
  * @param {T} key 要插入的键
  * @return {Node} 为了配合 insert 方法，一定要返回节点
  */
protected insertNode(node: Node<T>, key: T): Node<T> {
  // 与二叉搜索树的插入方式一致
  if (node == null) {
    return new Node(key);
  }
  if (this.compareFn(key, node.key) === Compare.LESS_THAN) {
    node.left = this.insertNode(node.left, key);
  } else if (this.compareFn(key, node.key) === Compare.BIGGER_THAN) {
    node.right = this.insertNode(node.right, key);
  } else {
    return node; // 重复的 key
  }
  // 校验树是否平衡
  return this.keepBalance(node);
}

/**
  * @description:  删除节点的递归方法，递归完成后也需要再平衡
  * @param {Node} node 要从中删除的节点
  * @param {T} key 要删除的键
  * @return {Node} 同样为了配合remove方法，一定要返回节点
  */
protected removeNode(node: Node<T>, key: T): Node<T> {
  // 与二叉搜索树的删除方式一致
  if (node == null) {
    return null;
  }

  if (this.compareFn(key, node.key) === Compare.LESS_THAN) {
    node.left = this.removeNode(node.left, key);
  } else if (this.compareFn(key, node.key) === Compare.BIGGER_THAN) {
    node.right = this.removeNode(node.right, key);
  } else {
    if (node.left == null && node.right == null) {
      node = null;
    } else if (node.left == null && node.right != null) {
      node = node.right;
    } else if (node.left != null && node.right == null) {
      node = node.left;
    } else {
      const aux = this.minNode(node.right);
      node.key = aux.key;
      node.right = this.removeNode(node.right, aux.key);
    }
  }

  // 校验树是否平衡
  return this.keepBalance(node);
}

虽然 AVL 树已经实现了二叉搜索树的平衡，但是由于每次添加和删除节点时都会涉及到再平衡操作，所以会比较复杂费时。所以在工业级应用中，都会用另一种近似平衡树算法红黑树来实现，下一篇来分析这个大名鼎鼎的 红黑树。

---

前端记事本，不定期更新，欢迎关注！

• 微信公众号： 林景宜的记事本
• 博客：林景宜的记事本
• 掘金专栏：林景宜的记事本
• 知乎专栏： 林景宜的记事本
• Github: MageeLin

---