Problem description

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

1. Let left be the count of nodes in left subtree, if k = 1 + left, then root is the answer. If k > 1 + left, then the answer is in right subtree with modified k = k - 1 - left

Below is python implementation

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution:
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        self.ans = 0
        
        def dfs(root:TreeNode, k:int)->int:
            if root == None:
                return 0
            left = dfs(root.left, k)
            if k == 1 + left:
                self.ans = root.val
            right = dfs(root.right, k - left - 1)
            return 1 + left + right
        
        dfs(root, k)
        return self.ans