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.
Example 1:
Input: root = [3,1,4,null,2], k = 1
3
/ \
1 4
\
2
Output: 1
Example 2:
Input: root = [5,3,6,2,4,null,null,1], k = 3
5
/ \
3 6
/ \
2 4
/
1
Output: 3
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
Use InOrder traversal
Check left sub-tree first, check current count, then check right sub-tree
class Solution {
private var count = 0
func kthSmallest(_ root: TreeNode?, _ k: Int) -> Int {
if let res = helper(root, k) {
return res
}
return -1
}
func helper(_ root: TreeNode?, _ k: Int) -> Int? {
// base case
if root == nil { return nil }
// search in left subtree
let left = helper(root?.left, k)
// if k'th smallest is found in left subtree, return it
if (left != nil) { return left }
count += 1
// if current element is k'th smallest, return it
if (count == k) { return root?.val }
// else search in right subtree
return helper(root?.right, k)
}
}