230. Kth Smallest Element in a BST**

230. Kth Smallest Element in a BST**

https://leetcode.com/problems/kth-smallest-element-in-a-bst/

题目描述

Given a binary search tree, write a function kthSmallest to find the kthsmallest element in it.

Note:
You may assume k is always valid, 1 ≤ k ≤ BST's total elements.

Example 1:

Input: root = [3,1,4,null,2], k = 1
   3
  / \
 1   4
  \
   2
Output: 1

Example 2:

Input: root = [5,3,6,2,4,null,null,1], k = 3
       5
      / \
     3   6
    / \
   2   4
  /
 1
Output: 3

Follow up:
What if the BST is modified (insert/delete operations) often and you need to find the kth smallest frequently? How would you optimize the kthSmallest routine?

C++ 实现 1

使用一个全局的计数器 count, 如果在中序遍历时刚好等于 k, 那么返回这个结果. 此题还有迭代的解法, 以及 Follow up 可以思考, 这部分内容可以参考: https://leetcode.com/problems/kth-smallest-element-in-a-bst/solution/.

class Solution {
private:
    int count = 0, res;
    void inorder(TreeNode *root, int k) {
        if (!root) return;
        inorder(root->left, k);
        count += 1;
        if (count == k) {
            res = root->val;
            return;
        }
        inorder(root->right, k);
    }
public:
    int kthSmallest(TreeNode* root, int k) {
        inorder(root, k);
        return res;
    }
};

C++ 实现 2

这道题也可以使用迭代的方式来做. 我们知道 BST 的中序遍历是有序数组, 那么只需要在中序遍历的时候将第 K 个元素输出即可. 对于二叉树中序遍历的迭代写法, 在 94. Binary Tree Inorder Traversal** 进行了介绍, 因此我们写出如下代码:

class Solution {
public:
    int kthSmallest(TreeNode* root, int k) {
        stack<TreeNode*> st;
        TreeNode *p = root;
        while (!st.empty() || p) {
            if (p) {
                st.push(p);
                p = p->left;
            } else {
                p = st.top();
                st.pop();
                if (--k == 0) return p->val;
                p = p->right;
            }
        }
        return -1;
    }
};

仅在中序遍历的迭代写法中插入了 if (--k == 0) return p->val; 这一行.

另外有个有意思的代码可以再记录下, 这是段 错误 的代码, 但 debug 花了一些时间, 一时不察写错了, 所以啊, 太晚写代码就容易头晕 …

/*
* 注意注意!!!! 这是错误代码 !!!!
*/
class Solution {
public:
    int kthSmallest(TreeNode* root, int k) {
        stack<TreeNode*> st;
        TreeNode *p = root;
        while (!st.empty() || p) {
            if (p) {
                st.push(p);
                p = p->left;
            } else {
                auto p = st.top();
                st.pop();
                if (--k == 0) return p->val;
                p = p->right;
            }
        }
        return -1;
    }
};
/*
* 注意注意!!!! 这是错误代码 !!!!
*/

表面上看逻辑和前一段代码一样, 但得不到正确的结果, 比如可以使用:

[3,1,4,null,2]
3

树:
      3
     / \
    1   4
     \
      2

这个例子, 上面代码输出的不是 3 而是 -1. 那么 bug 在哪呢?

bug 在 auto p = st.top(); 这一行.

我…用了个局部变量 auto p 覆盖了在 while 语句外面定义的 TreeNode *p …

而 auto p 只会在 else 逻辑块中生效. 当执行到 else 语句时, 全局的 TreeNode *p 此时肯定是空指针 (否则也不会执行到 else 语句), 而 else 语句会不断地将 stack 中元素弹出, 同时不会更新 TreeNode *p, 即 TreeNode *p 始终为空指针, 当 stack 元素全部弹出后, 由于 -- k == 0 最终没有满足, 所以最后执行的是 return -1.

忧伤, 这点要注意啊.