LeetCode230——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?

实现：
class Solution {
public:
    int kthSmallest(TreeNode* root, int k) {
        if (root == NULL) return 0;
        int val;
        kthSmallest(root, k, val);
        return val;
    }
    int kthSmallest(TreeNode* root, int k, int& val) {
        if (root == NULL) return k;
        int nk = k;
        if (root->left) {
            nk = kthSmallest(root->left, nk, val);
            if (nk == 0) {
                return 0; 
            }
        }
        nk--;
        if (nk == 0) {
            val = root->val;
            return 0;
        }
        if (root->right) {
            nk = kthSmallest(root->right, nk, val);
            if (nk == 0) {
                return 0; 
            }
        }
        return nk;
    }
};