Kth Smallest Element in a BST — Day 35(Python)

Photo by Todd Quackenbush on Unsplash

Today’s question is another favorite among major tech interviewers. It tests your knowledge about BST, recursion. Let us jump into the problem without wasting much time.

230. Kth Smallest Element in a BST

Given a binary search tree, write a function kth Smallest to find the kth smallest element in it.

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 kth smallest routine?

Constraints:

• The number of elements of the BST is between 1 to 10^4.
• You may assume k is always valid, 1 ≤ k ≤ BST's total elements.

Before we jump into the problem, we need to understand a few features of the Binary Search Tree.

A binary search tree is a tree-like data structure with added features as follows.

1. The subtree to the left side of the root node will always be lesser than the root node.
2. The subtree to the right side of the root node will always be greater than the root node.
3. The left and right subtree of the root node must be a Binary Search Tree.

One of the easiest ways to solve this problem is by traversing through the tree, storing the values of the node in a list. Later, sort the list and return the (K-1)th element in the list. For traversing the tree, we can choose any of the algorithms. Few traversal algorithms are as follows:

1. Inorder
2. Preorder
3. Postorder

Let us look at the code when we use inorder traversal.

class KFinder:
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        self.output = []
        def inorder(root):
            if root.left:
                inorder(root.left)
            self.output.append(root.val)
            if root.right:
                inorder(root.right)
        inorder(root)
        return(sorted(self.output)[k-1])

Complexity analysis.

Time Complexity

Inorder traversal visits each node once. Hence for inorder traversal, the time complexity is O(N). But we need to notice that sorting takes O(NlogN). Therefore the overall time complexity is O(NlogN).

Space Complexity

In the inorder traversal, each node will be visited once and stored in an array which takes O(N). Hence space complexity is O(N).

Did we notice, the inorder traversal of BST results in a sorted array? We are duplicating our effort by sorting the already sorted list. Hence we can remove the line of code that sorts the array. Let us look into the code snippet.

class KFinder:
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        self.output = []
        def inorder(root):
            if root.left:
                inorder(root.left)
            self.output.append(root.val)
            if root.right:
                inorder(root.right)
        inorder(root)
        return(self.output)[k-1]

Complexity analysis.

Time Complexity

Inorder traversal visits each node once. Hence for inorder traversal, the time complexity is O(N). Therefore the overall time complexity is O(N).

Space Complexity

In the inorder traversal, each node will be visited once and stored in an array which takes O(N). Hence space complexity is O(N).

Do we have to traverse through the whole tree? Can we try to return the answer once we find the Kth element? Do we have to visit nodes that are higher than the Kth value?

Let us see how we can return the result when we find the Kth value in our BST.

class KFinder:
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        self.output = []
        def inorder(root):
            if root.left:
                inorder(root.left)
            if len(self.output) < k:
                self.output.append(root.val)
            else:
                return
            if root.right:
                inorder(root.right)
        
        inorder(root)
        return self.output[-1]

Complexity analysis.

Time Complexity

In the above code, we will be visiting the leftmost node in the tree and then move K steps from that point. To visit the leftmost node, we will be traversing through the height of the tree. Hence the time complexity is O(H+k), where H is the height of the tree.

Space Complexity

In the above code, we will be visiting the leftmost node in the tree and then move K steps from that point. Since the code uses recursion, an internal stack is maintained. Hence the space complexity is O(H+k), where H is the height of the tree.

Do we need to have a stack to store all the values until we reach K? Can we use just a variable instead of a list?

Sure we can use a variable to store the answer, but how do we keep track of the Kth position? We can use another variable whose counter starts when we reach the leftmost value, and when this counter is equal to K, we have reached out Kth smallest element in the BST. Let us look into the code snippet.

class KFinder:
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        self.l = 0
        self.output = 0
        def helper(root):
            if root.left:
                helper(root.left)
            self.l += 1
            if self.l == k:
                self.output = root.val
                return
            if self.l > k:
                return
            if root.right:
                helper(root.right)
        helper(root)
        return self.output

Complexity analysis.

Time Complexity

In the above code, we will be visiting the leftmost node in the tree and then move K steps from that point. To visit the leftmost node, we will be traversing through the height of the tree. Hence the time complexity is O(H+k), where H is the height of the tree.

Space Complexity

In the above code, we will be visiting the leftmost node in the tree and then move K steps from that point. Since the code uses recursion, an internal stack is maintained. Hence the space complexity is O(H+k), where H is the height of the tree.

I would love to hear your feedback about my posts. Do let me know if you have any comments or feedback.