LeetCode:Kth Smallest Element in a BST

Kth Smallest Element in a BST
Total Accepted: 49142 
Total Submissions: 128501 
Difficulty: Medium
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?
Hint:
Try to utilize the property of a BST.
What if you could modify the BST node's structure?
The optimal runtime complexity is O(height of BST).
Credits: Special thanks to @ts for adding this problem and creating all test cases.
Subscribe to see which companies asked this question
Hide Tags
 
Binary Search Tree
Hide Similar Problems
 
(M) Binary Tree Inorder Traversal
java code:

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
public class Solution {
    public int kthSmallest(TreeNode root, int k) {
        
        int count = countNode(root.left);
        
        if(k == count + 1)
            return root.val;
        else if(k < count + 1)
            return kthSmallest(root.left, k);
        else // k > count + 1
            return kthSmallest(root.right, k - count - 1);
            
    }
    
    //      :        
    int countNode(TreeNode root) {
        if(root == null) return 0;
        
        return countNode(root.left) + countNode(root.right) + 1;
    }
    
}