230. Kth Smallest Element in a BST

Given a binary search tree, write a function

kthSmallest

to find thekth 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?

水一发啊，水一发～～～

/**
* 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:
void dfs(TreeNode *root, int &ret, int &k) {
if(k == 0) return ;
if(root->left != NULL) {
dfs(root->left, ret, k);
}
if(--k == 0) {
ret = root->val;
return ;
}
if(root->right != NULL) {
dfs(root->right, ret, k);
}
return ;
}
int kthSmallest(TreeNode* root, int k) {
if(root == NULL) return 0;
int ans;
dfs(root, ans, k);
return ans;
}
};

Follow Up:

在每个二叉树节点中添加一个记录以当前节点为根的子树中节点的个数，进行插入删除操作时，需要在插入点和删除点至根节点的路径上的所有节点的计数域进行维护，

然后查找的时候由于二叉查找树的性质，根据这个计数域进行查找即可。