[LeetCode] 230、二叉搜索树中第K小的元素

题目描述

给定一个二叉搜索树，编写一个函数 kthSmallest 来查找其中第 k 个最小的元素。

说明：
你可以假设 k 总是有效的，1 ≤ k ≤ 二叉搜索树元素个数。

示例：

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

进阶：
如果二叉搜索树经常被修改（插入/删除操作）并且你需要频繁地查找第 k 小的值，你将如何优化 kthSmallest 函数？

参考代码

/*
struct TreeNode {
    int val;
    struct TreeNode *left;
    struct TreeNode *right;
    TreeNode(int x) :
            val(x), left(NULL), right(NULL) {
    }
};
*/
class Solution {
public:
    int kthSmallest(TreeNode* root, int k) {
        if(root == nullptr || k <= 0)
            return 0;
            
        return KthNode(root, k)->val;
    }
    
    TreeNode* KthNode(TreeNode* pRoot, int &k)  // 这里一定要用 引用！！
    {
        if(pRoot == nullptr || k <= 0)
            return nullptr;
        
        TreeNode* pAns = nullptr;
        pAns = KthNode(pRoot->left, k);
        
        if(pAns == nullptr){
            k--;
            if(k == 0)
                pAns = pRoot;
        }
        
        if(pAns == nullptr)
            pAns = KthNode(pRoot->right, k);
        
        return pAns;
    }

};

另解：（推荐）

class Solution {
public:
    int n=0;
    int res;
    int kthSmallest(TreeNode* root, int k) {
        //递归
        dfs(root,k);
        return res;
    }

    void dfs(TreeNode* root,int k){
          if(!root)return ;
          dfs(root->left,k);
          n++;
          if(n==k) res=root->val;
          dfs(root->right,k);
    }
};