Taro Logo

Kth Smallest Element in a BST

Medium
Google logo
Google
2 views
Topics:
TreesRecursion

Given the root of a binary search tree, and an integer k, return the kth smallest value (1-indexed) of all the values of the nodes in the tree.

For example:

Consider the following binary search tree:

  3
 / \
1   4
 \ 
  2

If k = 1, the 1st smallest element is 1. If k = 2, the 2nd smallest element is 2. If k = 3, the 3rd smallest element is 3. If k = 4, the 4th smallest element is 4.

Another example:

    5
   / \
  3   6
 / \
2   4

/ 1

If k = 1, the 1st smallest element is 1. If k = 3, the 3rd smallest element is 3.

Your solution should be efficient and handle edge cases gracefully. What is the time and space complexity of your proposed solution? How would you optimize if the tree is frequently modified with insert and delete operations?

Solution


Kth Smallest Element in a BST

Problem Description

Given the root of a binary search tree (BST) and an integer k, the task is to find the kth smallest value (1-indexed) among all the node values in the BST.

Brute Force Solution

A straightforward approach is to perform an inorder traversal of the BST, which yields the elements in sorted order. We can store these elements in a list and then return the element at index k-1.

Code (Python)

def kth_smallest_brute_force(root, k):
    elements = []

    def inorder_traversal(node):
        if not node:
            return
        inorder_traversal(node.left)
        elements.append(node.val)
        inorder_traversal(node.right)

    inorder_traversal(root)
    return elements[k - 1]

Time Complexity

O(N), where N is the number of nodes in the BST, as we need to traverse all nodes.

Space Complexity

O(N), as we store all the node values in a list.

Optimal Solution

We can optimize the space complexity by performing an inorder traversal and maintaining a counter to keep track of the number of nodes visited. When the counter equals k, we return the current node's value.

Code (Python)

class Solution:
    def kthSmallest(self, root, k):
        self.count = 0
        self.result = None

        def inorder(node):
            if not node:
                return

            inorder(node.left)
            self.count += 1
            if self.count == k:
                self.result = node.val
                return
            inorder(node.right)

        inorder(root)
        return self.result

Time Complexity

O(H + k), where H is the height of the BST. In the worst-case scenario (skewed tree), H can be N, resulting in O(N) time complexity. However, in a balanced BST, H is log(N), resulting in O(log(N) + k) time complexity. If k is small, this is better than O(N).

Space Complexity

O(H), where H is the height of the BST, due to the recursive call stack. In the worst-case scenario (skewed tree), H can be N, resulting in O(N) space complexity. In a balanced BST, H is log(N), resulting in O(log(N)) space complexity.

Edge Cases

  • Empty Tree: If the root is None, it should probably return None or raise an exception, depending on the specific requirements.
  • k Out of Range: If k is less than 1 or greater than the number of nodes in the tree, it should either return None or raise an exception.
  • Duplicate Values: The algorithm works correctly even if there are duplicate values in the BST.

Follow-up: Frequent Modifications to the BST

If the BST is modified frequently (insertions and deletions), and we need to find the kth smallest element frequently, we can augment the BST by storing the size of the subtree rooted at each node. This allows us to find the kth smallest element in O(H) time, where H is the height of the tree.

To maintain the subtree sizes during insertions and deletions, we need to update the sizes of the affected nodes, which can also be done in O(H) time.