JS AVL树（数据结构）- 笔记

Code: 

/**
 * AVL 树
 * @class
 */
class AVLTree {
    /**
     * @type {TreeNode}
     */
    #root;

    /**
     * @constructor
     */
    constructor() {
        this.#root = null;
    }

    /**
     * 获取节点高度
     * @param {TreeNode} node 
     */
    height(node) {
        // 空节点高度为-1
        return node === null ? -1 : node.height;
    }

    /**
     * 更新节点高度
     * @param {TreeNode} node 
     */
    #updateHeight(node) {
        // 当前节点的高度等于最高子树的高度 + 1
        node.height = Math.max(this.height(node.left), this.height(node.right)) + 1;
    }

    /**
     * 获取平衡因子
     * @param {TreeNode} node 
     * @returns 
     */
    getBalanceFactor(node) {
        // 空节点平衡因子为 0 
        if (node === null) {
            return 0;
        }
        // 节点平衡因子 = 左子树高度 - 右子树高度
        return this.height(node.left) - this.height(node.right);
    }

    /**
     * 右旋
     * @param {TreeNode} node 
     */
    #rightRotate(node) {
        const child = node.left;
        const grandChild = child.right;
        child.right = node;
        node.left = grandChild;
        this.#updateHeight(node);
        this.#updateHeight(child);

        return child;
    }

    /**
     * 左旋
     * @param {TreeNode} node 
     */
    #leftRotate(node) {
        const child = node.right;
        const grandChild = child.left;
        child.left = node;
        node.right = grandChild;
        this.#updateHeight(node);
        this.#updateHeight(child);

        return child;
    }

    /**
     * 旋转节点使其平衡
     * @param {TreeNode} node 
     */
    #rotate(node) {
        const factor = this.getBalanceFactor(node);
        // 左偏
        if (factor > 1) {
            // 直接右旋
            if (this.getBalanceFactor(node.left) >= 0) {
                return this.#rightRotate(node);
            }
            // 先左旋后右旋
            node.left = this.#leftRotate(node.left);
            return this.#rightRotate(node);
        }
        // 右偏
        if (factor < -1) {
            // 直接左旋
            if (this.getBalanceFactor(node.right) <= 0) {
                return this.#leftRotate(node.right);
            }
            // 先右旋后左旋
            node.right = this.#rightRotate(node.right);
            return this.#leftRotate(node);
        }

        return node;
    }

    /**
     * 辅助节点的插入
     * @param {TreeNode} node 
     * @param {number} val 
     */
    #insertHelper(node, val) {
        if (node === null) {
            return new TreeNode(val);
        }
        if (val === node.val) {
            return node;
        }
        if (val < node.val) {
            node.left = this.#insertHelper(node.left, val);
        }
        else if (val > node.val) {
            node.right = this.#insertHelper(node.right, val);
        }
        this.#updateHeight(node);
        node = this.#rotate(node);

        return node;
    }

    /**
     * 插入节点
     * @param {number} val 
     */
    insert(val) {
        this.#root = this.#insertHelper(this.#root, val);
    }

    /**
     * 辅助节点的删除
     * @param {TreeNode} node 
     * @param {number} val 
     */
    #removeHelper(node, val) {
        if (node === null) {
            return null;
        }
        if (val < node.val) {
            node.left = this.#removeHelper(node.left, val);
        }
        else if (val > node.val) {
            node.right = this.#removeHelper(node.right, val);
        }
        else {
            if (node.left === null || node.right === null) {
                const child = node.left === null ? node.right : node.left;
                if (child === null) {
                    return null;
                }
                node = child;
            }
            else {
                let tmp = node.right;
                while(tmp.left !== null) {
                    tmp = tmp.left;
                }
                node.right = this.#removeHelper(node.right, tmp.val);
                node.val = tmp.val;
            }
        }
        this.#updateHeight(node);
        node = this.#rotate(node);

        return node;
    }

    /**
     * 删除节点
     * @param {number} val 
     */
    remove(val) {
        this.#root = this.#removeHelper(this.#root, val);
    }

    /**
     * 查找节点
     * @param {number} val 
     * @returns 
     */
    search(val) {
        let cur = this.#root;
        while (cur !== null) {
            if (cur.val < val) {
                cur = cur.right;
            }
            else if (cur.val > val) {
                cur = cur.left;
            }
            else {
                break;
            }
        }

        return cur;
    }
}

/**
 * 树节点
 * @class
 */
class TreeNode {
    val;
    left;
    right;
    height;
    constructor(val, left, right) {
        this.val = val === undefined ? 0 : val;
        this.left = left === undefined ? left : null;
        this.right = right === undefined ? right : null;
        this.height = height === undefined ? 0 : height;
    }
}