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);

    }

}