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?

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

Similar: Binary Tree Inorder Traversal

 1 Solution 1. Binary Search
 2 Time Complexity: T(n) = n/2 + T(n/2) = O(n)
 3 /**
 4  * Definition for a binary tree node.
 5  * public class TreeNode {
 6  *     int val;
 7  *     TreeNode left;
 8  *     TreeNode right;
 9  *     TreeNode(int x) { val = x; }
10  * }
11  */
12 public class Solution {
13     public int kthSmallest(TreeNode root, int k) {
14         int count = countNodes(root.left);
15         if (k <= count) {
16             return kthSmallest(root.left, k);
17         } else if (k > count + 1) {
18             return kthSmallest(root.right, k-1-count); 
19         }
20         return root.val;
21     }
22 
23     public int countNodes(TreeNode n) {
24         if (n == null) return 0;
25         return 1 + countNodes(n.left) + countNodes(n.right);
26     }
27 }
28 
29 Solution 2. Modify the structure of BST
30 Time Complexity: O(lg n)
31 /**
32  * Definition for a binary tree node.
33  * public class TreeNode {
34  *     int val;
35  *     TreeNode left;
36  *     TreeNode right;
37  *     TreeNode(int x) { val = x; }
38  * }
39  */
40 public class Solution {
41     public int kthSmallest(TreeNode root, int k) {
42         TreeWithCount twcRoot = buildTWC(root);
43         return findKth(twcRoot, k);
44     }
45     
46     private int findKth(TreeWithCount root, int k) {
47         if (k <= 0 || k > root.count) return -1;
48         if (root.left != null) {
49             if (root.left.count == k-1) return root.val;
50             if (k <= root.left.count) return findKth(root.left, k);
51             return findKth(root.right, k-1-root.left.count);
52         } else {
53             if (k == 1) return root.val;
54             return findKth(root.right, k-1);
55         }
56     }
57     
58     private TreeWithCount buildTWC(TreeNode root) {
59         if (root == null) return null;
60         
61         TreeWithCount twcRoot = new TreeWithCount(root.val);
62         twcRoot.left = buildTWC(root.left);
63         twcRoot.right = buildTWC(root.right);
64         if (twcRoot.left != null) twcRoot.count += twcRoot.left.count;
65         if (twcRoot.right != null) twcRoot.count += twcRoot.right.count;
66         return twcRoot;
67     }
68     
69     class TreeWithCount {
70         int val;
71         int count;
72         TreeWithCount left;
73         TreeWithCount right;
74         public TreeWithCount (int val) {
75             this.val = val;
76             this.count = 1;
77         }
78     } 
79 }