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?

最先想到的思路就是中序遍历后找到结果数组中的第k-1项，为了节约时间可以每次都判断结果数组长度，达到k即返回最后一项。

def kth_smallest(root, k)
    stack,ans = [[false,root]], []
    while not stack.empty?
        visit, node = stack.pop
        if node
            if visit
               ans << node.val
               return ans[-1] if ans.length == k
            else
                stack << [false,node.right] << [true,node] << [false,node.left]
            end
        end
    end
end

 

Try the left subtree first. If that made k zero, then its answer is the overall answer and we return it right away. Otherwise, decrease k for the current node, and if that made k zero, then we return the current node's value right away. Otherwise try the right subtree and return whatever comes back from there.

int kthSmallest(TreeNode* root, int k) {
    return find(root, k);
}
int find(TreeNode* root, int& k) {
    if (root) {
        int x = find(root->left, k);
        return !k ? x : !--k ? root->val : find(root->right, k);
    }
}