May Leetcode Challenge

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

}