[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:
Try to utilize the property of a BST.
What if you could modify the BST node's structure?
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小的节点。