「这是我参与2022首次更文挑战的第18天，活动详情查看：2022首次更文挑战」

leetcode-230 二叉搜索树中第K小的元素

题目介绍

给定一个二叉搜索树的根节点 root ，和一个整数 k ，请你设计一个算法查找其中第 k 个最小元素（从 1 开始计数）。

示例1

输入： root = [3,1,4,null,2], k = 1
输出： 1

示例2

输入： root = [5,3,6,2,4,null,null,1], k = 3
输出： 3

提示：

• 树中的节点数为 n 。
• 1 <= k <= n <= 104
• 0 <= Node.val <= 104 进阶： 如果二叉搜索树经常被修改（插入/删除操作）并且你需要频繁地查找第 k 小的值，你将如何优化算法？

解题思路

此题根据二叉搜索树和中序遍历的性质可以知道，对二叉搜索树进行遍历，将会得到一个递增的数组，想要获取二叉搜索树中的第 K 小元素，只需要返回数组的第 K 个值就可以了

思路一：中序遍历-递归

解题步骤

1. 定义 ans 数组存放中序遍历之后的数组
2. 从根节点开始，递归遍历左子树，然后将根节点插入数组，再递归遍历右子树
3. 返回 ans 数组中下标为 k - 1 的元素

解题代码

var kthSmallest = function(root, k) {
    const ans = []
    inOrder(root, ans)
    return ans[k - 1]
};

var inOrder = function(root, ans) {
    if (!root) return
    inOrder(root.left, ans)
    ans.push(root.val)
    inOrder(root.right, ans)
}

思路二：中序遍历-迭代

迭代相比于递归是可以控制遍历的过程，当迭代到二叉搜索树的第 K 个节点时，就可以结束二叉搜索树的中序遍历了，不需要遍历完整颗二叉搜索树

解题代码

var kthSmallest = function(root, k) {
    // nodeArr 用于临时保存节点
    const nodeArr = []
    // 当遍历到第 K 个元素 或者 遍历完整棵二叉搜索树时，结束
    while (k-- && (nodeArr.length || root)) {
        // 当 root 不为空时，将 root 插入到 nodeArr 中，然后继续判断其左节点
        while (root) {
            nodeArr.push(root)
            root = root.left
        }
        // 从 nodeArr 数组中弹出一个元素，弹出的顺序即为中序遍历的顺序
        root = nodeArr.pop()
        k && (root = root.right)
    }
    return root.val
};

思路三：记录树的节点数

这种方法是利用了二叉搜索树的性质，如果值小于 根节点，那么将会出现在 根节点 的左子树中，如果值大于 根节点，将会出现在 根节点 的右子树中

因此，可以统计 根节点 左子树的节点数量

• 如果左子树的节点数量小于 k - 1，说明第 k 小的元素出现在根节点的右子树中
• 如果左子树的节点数量等于 k - 1，说明第 k 小的元素就是根节点
• 如果左子树的节点数量大于 k - 1，说明第 k 小的元素出现在根节点的左子树中

解题步骤

1. 定义一个 BST 类，用于统计并记录以每个节点为根节点的子树的节点数量
2. 根据以当前节点为根节点的左子树的数量，判断是否继续查找左子树或者右子树
3. 最后返回第 k 小的元素

解题代码

var kthSmallest = function(root, k) {
    const myBst = new BST(root)
    return myBst.getKth(k)
};

class BST {
    constructor(root) {
        this.root = root
        this.map = new Map()
        this.nodeNums(root)
    }

    // 用于记录以每个节点为根节点的子树的节点数量
    // 每棵子树的节点数量等于 左子树的节点数量 + 右子树的节点数量 + 1（根节点）
    nodeNums(root) {
        if (!root) return 0
        this.map.set(root, this.nodeNums(root.left) + this.nodeNums(root.right) + 1)
        return this.map.get(root)
    }

    // 获取当前子树的节点数量
    getNums(root) {
        return this.map.get(root) || 0
    }

    // 获取第 k 小的元素
    getKth(k) {
        let root = this.root
        while (root) {
            // 获取左子树的节点数量
            const left = this.getNums(root.left)
            // 如果左子树的节点数量小于 k-1，则从右子树中查找第 k-left-1 小的元素
            if (left < k - 1) {
                root = root.right
                k = k - left - 1
            } else if (left === k - 1) {
            // 如果左子树的节点数量等于 k-1，说明根节点为第 k 小的元素，直接跳出循环
                break
            } else {
            // 如果左子树的节点数量大于 k-1，则从根节点的左子树中继续查找第 k 小的元素
                root = root.left
            }
        }
        // 最后返回 root 节点的值
        return root.val
    }
}