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

cpp 复制代码

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
class Solution {
public:
    int kthSmallest(TreeNode* root, int k) {
        stack<TreeNode *>stk;
        while(root!=nullptr || stk.empty()==false){
            while(root!=nullptr){
                stk.push(root);
                root = root->left;
            }
            root = stk.top();
            stk.pop();
            k--;
            if(k==0){
                break;
            }
            root = root->right;
        }
        return root->val;
    }
};

这题本质上就是利用 二叉搜索树 BST 的中序遍历结果是递增序列。

你的核心思路是：

复制代码

BST 中序遍历：
左 → 根 → 右

得到的节点值：
从小到大排列

所以我们不需要真的把所有节点存进数组，只要在中序遍历过程中，每访问一个节点就：

复制代码

k--;

当：

复制代码

k == 0

时，当前节点就是第 k 小的节点。

不过你这份代码有一个小问题：

复制代码

stack<TreeNode *> stk;
stk.push(root);

这里一开始不应该先把 root 压栈，因为下面内层循环又会：

复制代码

while(root != nullptr){
    stk.push(root);
    root = root->left;
}

这样根节点会被压两次。

建议改成：

复制代码

class Solution {
public:
    int kthSmallest(TreeNode* root, int k) {
        stack<TreeNode*> stk;

        while (root != nullptr || !stk.empty()) {

            // 1. 一直向左走，把节点压栈
            while (root != nullptr) {
                stk.push(root);
                root = root->left;
            }

            // 2. 左边到底，弹出当前节点
            root = stk.top();
            stk.pop();

            // 3. 中序遍历访问到一个节点
            k--;

            // 第 k 小
            if (k == 0) {
                return root->val;
            }

            // 4. 处理右子树
            root = root->right;
        }

        return -1;
    }
};

你可以把整个框架记成：

复制代码

一直压左
↓
弹出根节点
↓
k--
↓
k == 0 就返回
↓
转向右子树

为什么这样能找到第 k 小？

例如 BST：

复制代码

中序遍历：

复制代码

1 → 2 → 3 → 4 → 5 → 6

假设：

复制代码

k = 3

那么访问：

复制代码

1   k: 3 → 2
2   k: 2 → 1
3   k: 1 → 0

所以答案就是：

复制代码

易错点

最重要的第一个易错点就是你代码里的：

复制代码

stk.push(root);

不要提前压栈。

因为标准中序遍历已经会在：

复制代码

while(root != nullptr)

里面把根节点压进去。

第二个易错点是外层循环：

复制代码

while(root != nullptr || !stk.empty())

不能只写：

复制代码

while(root != nullptr)

因为当左边走到底：

复制代码

root == nullptr

时，栈里面其实还有很多父节点没有访问。

第三个易错点是 k-- 的位置。

必须在：

复制代码

root = stk.top();
stk.pop();

之后执行：

复制代码

k--;

因为弹出节点的这一刻才是真正的"中序访问节点"。

不能在压栈的时候 k--，因为：

复制代码

stk.push(root);

只是保存节点，并不代表已经按照中序顺序访问了它。

第四个易错点是：

复制代码

root = root->right;

不能漏。

因为中序是：

复制代码

左 → 根 → 右

访问完当前根节点之后，必须去处理它的右子树。

最后你可以把这道题直接背成一句模板：

BST + 第 k 小 = 中序遍历，因为 BST 中序遍历天然从小到大；每访问一个节点就 k--，k == 0 时返回当前节点。

复杂度方面，最坏情况下时间复杂度是 O(n) ，栈的空间复杂度是 O(h) ，其中 h 是树的高度。