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

题目链接 · 灵神原题解（署名来源）

前置题目

98. 验证二叉搜索树，请看视频讲解【基础算法精讲 11】。

思路

由于中序遍历就是在从小到大遍历节点值，所以遍历到的第 $k$ 个节点值就是答案。

写法一：记录答案

在中序遍历，即「左-根-右」的过程中，每次递归完左子树，就把 $k$ 减少 $1$ ，表示我们按照中序遍历访问到了一个节点。如果减一后 $k$ 变成 $0$ ，那么答案就是当前节点的值，用一个外部变量 $ans$ 记录。

Python 解法1C++ 解法1

python

class Solution:
    def kthSmallest(self, root: Optional[TreeNode], k: int) -> int:
        ans = 0
        def dfs(node: Optional[TreeNode]) -> None:
            nonlocal k, ans
            if node is None or k <= 0:
                return
            dfs(node.left)  # 左
            k -= 1
            if k == 0:
                ans = node.val  # 根
            dfs(node.right)  # 右
        dfs(root)
        return ans

cpp

// C++ 版待补充

C++ 解法2

cpp

class Solution {
public:
    int kthSmallest(TreeNode* root, int k) {
        int ans;
        auto dfs = [&](this auto&& dfs, TreeNode* node) -> void {
            if (node == nullptr || k <= 0) {
                return;
            }
            dfs(node->left); // 左
            if (--k == 0) {
                ans = node->val; // 根
            }
            dfs(node->right); // 右
        };
        dfs(root);
        return ans;
    }
};

复杂度分析

时间复杂度： $O (n)$ ，其中 $n$ 是二叉树的大小（节点个数）。
空间复杂度： $O (h)$ ，其中 $h$ 是树高，递归需要 $O (h)$ 的栈空间。最坏情况下树是一条链， $h = n$ ，空间复杂度为 $O (n)$ 。

写法二：不记录答案 + 提前返回

写法一使用了一个外部变量记录答案，能否不使用外部变量记录呢？

可以，做法如下：

递归边界：如果当前节点是空节点，返回 $- 1$ ，表示没有找到。注意题目保证节点值非负。
执行中序遍历，先递归左子树。
判断左子树的返回值 $leftRes$ 是否为 $- 1$ 。如果不是 $- 1$ ，说明我们在左子树中找到了答案，返回 $leftRes$ 。如果是 $- 1$ ，说明尚未找到答案，继续下一步。
把 $k$ 减少 $1$ 。如果 $k = 0$ ，那么答案就是当前节点值，返回当前节点值。
现在，答案要么在当前节点的右子树中，要么在除了当前子树的其余节点中。递归右子树，如果答案在右子树中，那么直接返回答案；如果答案不在右子树中，那么右子树也会返回 $- 1$ ，由于当前子树搜索完毕，所以当前子树没有找到答案，返回 $- 1$ 。综上所述，可以直接返回右子树的返回值。

Python 解法3C++ 解法3

python

class Solution:
    def kthSmallest(self, root: Optional[TreeNode], k: int) -> int:
        def dfs(node: Optional[TreeNode]) -> int:
            if node is None:
                return -1  # 题目保证节点值非负，用 -1 表示没有找到
            left_res = dfs(node.left)
            if left_res != -1:  # 答案在左子树中
                return left_res
            nonlocal k
            k -= 1
            if k == 0:  # 答案就是当前节点
                return node.val
            return dfs(node.right)  # 右子树会返回答案或者 -1
        return dfs(root)

cpp

// C++ 版待补充

C++ 解法4

cpp

class Solution {
public:
    int kthSmallest(TreeNode* root, int& k) { // 注意这里改成了引用
        if (root == nullptr) {
            return -1; // 题目保证节点值非负，用 -1 表示没有找到
        }
        int left_res = kthSmallest(root->left, k);
        if (left_res != -1) { // 答案在左子树中
            return left_res;
        }
        if (--k == 0) { // 答案就是当前节点
            return root->val;
        }
        return kthSmallest(root->right, k); // 右子树会返回答案或者 -1
    }
};

C++ 解法5

cpp

class Solution {
public:
    int kthSmallest(TreeNode* root, int k) {
        auto dfs = [&](this auto&& dfs, TreeNode* node) -> int {
            if (node == nullptr) {
                return -1; // 题目保证节点值非负，用 -1 表示没有找到
            }
            int left_res = dfs(node->left);
            if (left_res != -1) { // 答案在左子树中
                return left_res;
            }
            if (--k == 0) { // 答案就是当前节点
                return node->val;
            }
            return dfs(node->right); // 右子树会返回答案或者 -1
        };
        return dfs(root);
    }
};

复杂度分析

时间复杂度： $O (n)$ ，其中 $n$ 是二叉树的大小（节点个数）。
空间复杂度： $O (h)$ ，其中 $h$ 是树高，递归需要 $O (h)$ 的栈空间。最坏情况下树是一条链， $h = n$ ，空间复杂度为 $O (n)$ 。

思考题

改成求第 $k$ 大要怎么做？

更多相似题目，见下面链表与二叉树题单中的「二叉搜索树」。

分类题单

如何科学刷题？

我的题解精选（已分类）

欢迎关注 B站@灵茶山艾府

前置题目 ​

思路 ​

写法一：记录答案 ​

复杂度分析 ​

写法二：不记录答案 + 提前返回 ​

复杂度分析 ​

思考题 ​

分类题单 ​

前置题目

思路

写法一：记录答案

复杂度分析

写法二：不记录答案 + 提前返回

复杂度分析

思考题

分类题单