230. Kth Smallest Element in a BST#

Difficulty: Medium

Link to Problem: To see the Kth Smallest Element in a BST problem on LeetCode, click here!

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

Constraints:

  • The number of nodes in the tree is n.

  • 1 <= k <= n <= \(10^4\)

  • 0 <= Node.val <= \(10^4\)

Follow up: If the BST is modified often (i.e., we can do insert and delete operations) and you need to find the kth smallest frequently, how would you optimize?

Solution:#

Here’s a Python function to solve this problem:

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

def kthSmallest(root, k):
    stack = []  # Initialize an empty stack to simulate the traversal
    while root or stack:
        while root:
            stack.append(root)  # Push the current node onto the stack
            root = root.left  # Move to the left child
        root = stack.pop()  # Pop a node from the stack
        k -= 1  # Decrement k since we've visited a node
        if k == 0:
            return root.val  # If k becomes 0, return the current node's value as the kth smallest
        root = root.right  # Move to the right child to continue the traversal

Explanation:#

  1. We start by defining a TreeNode class, which represents a node in the binary search tree (BST). Each node has a value (val), a left child (left), and a right child (right).

  2. The kthSmallest function takes two parameters: root, which is the root of the BST, and k, which represents the kth smallest element we want to find.

  3. We initialize an empty stack (stack) to simulate the traversal of the BST. The stack will be used to keep track of nodes as we traverse the tree in an iterative manner.

  4. We enter a while loop that continues until either root becomes None (indicating we have traversed the entire tree) or the stack is empty.

  5. Within the loop, we start another while loop to traverse as far left as possible in the BST. We repeatedly push nodes onto the stack and move to their left children until we reach the leftmost leaf node. This is the smallest node in the BST.

  6. Once we have reached the leftmost leaf node (the smallest node), we pop nodes from the stack one by one. As we pop each node, we decrement k by 1 to keep track of the number of nodes we have visited.

  7. If k becomes 0 after decrementing, it means we have found the kth smallest element. In this case, we return the val of the current node as the result.

  8. If k is still greater than 0, it means we haven’t found the kth smallest element yet. In this case, we move to the right child of the current node to continue the traversal, as the kth smallest element, if it exists, will be in the right subtree of the current node.

  9. The process continues until we find the kth smallest element or traverse the entire tree.

  10. Finally, we return the kth smallest element found.

The key idea in this code is to perform an in-order traversal of the BST while keeping track of the kth smallest element. By visiting nodes in ascending order, we can efficiently find the kth smallest element in O(h + k) time, where h is the height of the BST and k is the desired kth element. This approach is both concise and efficient for this problem.

Test cases#

# Example 1: 

# Example usage:
# Create the tree for the first example: root = [3,1,4,null,2], k = 1
root1 = TreeNode(3)
root1.left = TreeNode(1)
root1.right = TreeNode(4)
root1.left.right = TreeNode(2)
k1 = 1
print(kthSmallest(root1, k1)) 

1

# Example 2:

# Create the tree for the second example: root = [5,3,6,2,4,null,null,1], k = 3
root2 = TreeNode(5)
root2.left = TreeNode(3)
root2.right = TreeNode(6)
root2.left.left = TreeNode(2)
root2.left.right = TreeNode(4)
root2.left.left.left = TreeNode(1)
k2 = 3
print(kthSmallest(root2, k2))

3

Time and Space Complexity Analysis#

Let’s analyze the time and space complexity of the given code:

Time Complexity:

The time complexity of this code is O(h + k), where:

  • h is the height of the binary search tree (BST).

  • k is the desired kth smallest element we want to find.

  1. In the worst case, where the BST is highly unbalanced and resembles a linked list, the height (h) of the tree can be equal to the number of nodes (n) in the tree. In this case, the time complexity is O(n + k).

  2. In the best case, where the BST is perfectly balanced, the height (h) is log(n), where n is the number of nodes in the tree. In this case, the time complexity is O(log(n) + k).

So, the time complexity can vary from O(log(n) + k) in the best-case scenario to O(n + k) in the worst-case scenario. Typically, for balanced BSTs, the time complexity is closer to O(log(n) + k).

Space Complexity:

The space complexity of this code is O(h) due to the stack used for the iterative in-order traversal, where h is the height of the BST.

  1. In the worst case, when the BST is highly unbalanced and resembles a linked list, the height (h) of the tree can be equal to the number of nodes (n) in the tree. In this case, the space complexity is O(n) because the stack can potentially store all n nodes.

  2. In the best case, when the BST is perfectly balanced, the height (h) is log(n), where n is the number of nodes in the tree. In this case, the space complexity is O(log(n)) because the stack will have at most log(n) nodes.

So, the space complexity depends on the height of the BST and can vary from O(log(n)) in the best-case scenario to O(n) in the worst-case scenario.

In practical terms, for balanced BSTs or moderately unbalanced BSTs, the space complexity is usually close to O(log(n)), and the code is efficient in terms of space usage.

Challenging Exercises:#

  1. Kth Largest Element: Modify the code to find the kth largest element in the BST instead of the kth smallest element.

  2. Kth Smallest Element in Two BSTs: Given two BSTs, find the kth smallest element when considering both BSTs as a single sorted list. This involves merging the two BSTs while finding the kth element efficiently.