Question
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?
My Solution
/**
* Definition for a binary tree node.
* struct TreeNode {
* int val;
* TreeNode *left;
* TreeNode *right;
* TreeNode(int x) : val(x), left(NULL), right(NULL) {}
* };
*/
class Solution {
public:
int kthSmallest(TreeNode* root, int k) {
/*中序遍历,取最小的k个*/
int rst, num;
rst = middleIterate(root, k, num);
return rst;
}
int middleIterate(TreeNode* root, int k, int& num){
/* 中序遍历树
0. 边界条件:if root == Null then num = 0, return 0;
1. rst = middleIterate(root->left, k, lnum)遍历左子树
2. if lnum >= k then num = lnum, return rst;
3. else if lnum == k - 1 then num = k, return root->val;
4. else k = k - lnum - 1, rst = middleIterate(root->left, k, rnum)遍历右子树
5. num = lnum + rnum + 1, return rst;
*/
if(NULL == root){
num = 0;
return 0;
}
int r, lnum, rnum;
r = middleIterate(root->left, k, lnum);
if(lnum >= k){
num = lnum;
return r;
}else if(k - 1 == lnum)
{
num = k;
return root->val;
}else{
k = k - lnum - 1;
r = middleIterate(root->right, k, rnum);
num = lnum + rnum + 1;
return r;
}
}
};