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?

Solution:

The easy solution would be in order traverse the tree and return the kth element in order. 

Please refer to this post to see all three in order traversal techniques.: http://codesniper.blogspot.com/2015/03/94-binary-tree-inorder-traversal.html

  public int kthSmallest(TreeNode root, int k) {  
     Stack<TreeNode> st=new Stack<TreeNode>();  
     int num=0;  
     while(root!=null || !st.empty()){  
       if(root!=null){  
         st.push(root);  
         root=root.left;  
       }  
       else{  
         root=st.pop();  
         num++;  
         if(num==k) return root.val;  
         root=root.right;  
       }  
     }  
     return 0;  
   }  

To the follow up question:
What if we modify the BST often and need to find Kth smallest number often.
We should create a TreeNodeWithCount object which extends the basic treenode object but with a property of the counts of this root node. In this way, once we have this treeNodeWithCount initiated(O(n)), any insert/delete operation will take O(lg(n) to update the tree including the counts and we can find Kth smallest number in O(lg(n)).