js实现平衡二叉树

个人学习记录。

// 定义节点类
class TreeNode {
  constructor(key, value) {
    this.key = key;
    this.value = value;
    this.left = null;
    this.right = null;
    this.height = 1; // 初始高度为1
  }
}

// 定义平衡二叉搜索树类
class AVLTree {
  constructor() {
    this.root = null;
  }

  // 获取节点高度
  getHeight(node) {
    return node ? node.height : 0;
  }

  // 计算节点平衡因子（左子树高度 - 右子树高度）
  getBalanceFactor(node) {
    return node ? this.getHeight(node.left) - this.getHeight(node.right) : 0;
  }

  // 左旋操作
  rotateLeft(node) {
    const rightChild = node.right;
    node.right = rightChild.left;
    rightChild.left = node;

    node.height = Math.max(this.getHeight(node.left), this.getHeight(node.right)) + 1;
    rightChild.height = Math.max(this.getHeight(rightChild.left), this.getHeight(rightChild.right)) + 1;

    return rightChild;
  }

  // 右旋操作
  rotateRight(node) {
    const leftChild = node.left;
    node.left = leftChild.right;
    leftChild.right = node;

    node.height = Math.max(this.getHeight(node.left), this.getHeight(node.right)) + 1;
    leftChild.height = Math.max(this.getHeight(leftChild.left), this.getHeight(leftChild.right)) + 1;

    return leftChild;
  }

  // 更新节点高度并检查是否需要进行旋转操作以保持平衡
  balance(node) {
    node.height = Math.max(this.getHeight(node.left), this.getHeight(node.right)) + 1;

    const balanceFactor = this.getBalanceFactor(node);

    if (balanceFactor > 1) { // 左子树过高
      if (this.getBalanceFactor(node.left) < 0) {
        node.left = this.rotateLeft(node.left);
      }
      return this.rotateRight(node);
    } else if (balanceFactor < -1) { // 右子树过高
      if (this.getBalanceFactor(node.right) > 0) {
        node.right = this.rotateRight(node.right);
      }
      return this.rotateLeft(node);
    }

    return node;
  }

  // 插入节点
  insert(key, value) {
    this.root = this._insertNode(this.root, key, value);
  }

  _insertNode(node, key, value) {
    if (!node) {
      return new TreeNode(key, value);
    }

    if (key < node.key) {
      node.left = this._insertNode(node.left, key, value);
    } else if (key > node.key) {
      node.right = this._insertNode(node.right, key, value);
    } else {
      node.value = value; // 如果键已存在，则更新值
    }

    return this.balance(node);
  }

  // 为了完整演示，这里提供一个简单的查找方法
  find(key) {
    return this._findNode(this.root, key);
  }

  _findNode(node, key) {
    if (!node || node.key === key) {
      return node;
    }

    return key < node.key ? this._findNode(node.left, key) : this._findNode(node.right, key);
  }
}

// 使用示例
const tree = new AVLTree();
tree.insert(5, 'value5');
tree.insert(3, 'value3');
tree.insert(7, 'value7');
tree.insert(2, 'value2');
tree.insert(4, 'value4');
tree.insert(6, 'value6');
tree.insert(8, 'value8');

console.log(tree.find(5).value); // 输出: "value5"