44.二叉搜索树中第 K 小的元素

题目链接

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

解法1 暴力解法

思路

根据二叉搜索树的特性，中序遍历出来的二叉搜索树是严格递增，于是我们可以在中序遍历的过程中去保存节点的值。最后通过找寻下标的方式去返回结果。

代码

function kthSmallest(root: TreeNode | null, k: number): number {
    const list = [];

    const inorder = (node) => {
        if (!node) return;
        inorder(node.left);
        list.push(node.val);
        inorder(node.right);
    };

    inorder(root);

    return list[k - 1];
};

时空复杂度

时间复杂度：O(n)

空间复杂度：缓存了所有节点 O(n)

解法2 中序+计数

思路

在刚刚的暴力解法当中，我们缓存了所有的节点，可不可以不保存下来呢？

二叉搜索树严格递增，在中序遍历时去缓存一个变量记录当前是第几个节点，如果当前节点恰好是第 k 小的节点，将答案更新即可。最后返回答案。

代码

function kthSmallest(root: TreeNode | null, k: number): number {
    let result = -1;
    let counter = 0;

    const inorder = (node) => {
        if (!node) return;
        inorder(node.left);

        counter++; // 先自增再比较，因为k是从1计数
        if (counter === k) {
            result = node.val;
            return;
        }

        inorder(node.right);
    };

    inorder(root);

    return result;
};

不使用递归的栈版中序遍历

function kthSmallest(root: TreeNode | null, k: number): number {
    const stack: TreeNode[] = [];

    let current = root;
    while (current || stack.length > 0) {
        // 一直往左走
        while (current) {
            stack.push(current);
            current = current.left;
        }

        current = stack.pop()!;
        k--;

        if (k === 0) {
            return current.val;
        }

        current = current.right;
    }

    return -1; // 正常不会到这里
};

时空复杂度

时间复杂度：最好 O(k)，最差 O(n)

空间复杂度：递归栈空间开销 O(logn)