[leetcode] 230. Kth Smallest Element in a BST

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

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

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?

Hint:

1. Try to utilize the property of a BST.
2. What if you could modify the BST node's structure?
3. The optimal runtime complexity is O(height of BST).

这道题是找BST中值第K小的节点，题目难度为Medium。

中序遍历BST得到的节点值是递增的，所以题目可以转化为中序遍历二叉树获取第K个节点，第98题（传送门）和这道题类似，大家可以先看下第98题。

中序遍历二叉树就不详细介绍了，不清楚的同学可以看下第94题（传送门）。这里列出递归和非递归两个版本的代码。

递归版本代码：

class Solution {
    void getKthSmallest(TreeNode* root, int& k, int& val) {
        if(!root) return;
        if(root->left) getKthSmallest(root->left, k, val);
        --k;
        if(!k) {
            val = root->val;
            return;
        }
        if(root->right) getKthSmallest(root->right, k, val);
    }
public:
    int kthSmallest(TreeNode* root, int k) {
        int ret;
        getKthSmallest(root, k, ret);
        return ret;
    }
};

非递归版本代码：

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

另外还可以通过二分查找法确定值第K小的节点，如果当前节点左子树节点数为K-1，第K小的节点即是当前节点；如果左子树节点数大于或等于K，表明第K小节点在左子树，继续在左子树中二分查找；如果左子树节点数小于K-1，表明第K小节点在右子树，继续在右子树中二分查找。具体代码：

class Solution {
    int getNodesCnt(TreeNode* n) {
        if(!n) return 0;
        return getNodesCnt(n->left) + getNodesCnt(n->right) + 1;
    }
public:
    int kthSmallest(TreeNode* root, int k) {
        if(!root) return 0;
        int cnt = getNodesCnt(root->left);
        if(cnt >= k)
            return kthSmallest(root->left, k);
        else if(cnt < k - 1)
            return kthSmallest(root->right, k-cnt-1);
        else 
            return root->val;
    }
};

题目追问如果可以改变节点数据结构如何优化？可以在节点中加入子树节点个数，这样按照上面二分查找的方法即可在O(height of BST)时间复杂度内找出值第K小的节点。