Kth Smallest/Largest in BST

The Core Idea: Traversal is Sorting

The problem asks for the k-th smallest element. Given what we know from the last lesson, if we perform an in-order traversal (Left -> Node -> Right), we are essentially "visiting" the nodes in their sorted order. Therefore, the k-th smallest element is simply the k-th node we visit during an in-order traversal.

A naive approach would be to perform a full in-order traversal, store all the node values in a list, and then return the element at index k-1.

# Naive approach
def kthSmallest_naive(root, k):
    in_order_list = []
    
    def traverse(node):
        if not node:
            return
        traverse(node.left)
        in_order_list.append(node.val)
        traverse(node.right)
    
    traverse(root)
    return in_order_list[k-1]

This works, but it has a time complexity of $O(N)$ and, more importantly, a space complexity of $O(N)$ to store the list. In an interview, you'll be expected to do better, especially on space. The key is to stop the traversal as soon as we find the k-th element.

Approach 1: Recursive In-order Traversal with a Counter

We can adapt the recursive in-order traversal to stop early. The idea is to maintain a counter or decrement k each time we "visit" a node (after processing its left subtree). When the count reaches k (or k reaches 0), we've found our target.

This guide provides a very clear plan and implementation for this recursive strategy.

The core logic, as shown in the article, involves a helper function that passes k by reference (or uses nonlocal in Python) and decrements it at the "visit" step of the traversal. This stops unnecessary work and keeps space complexity tied to the recursion depth, $O(H)$, where $H$ is the tree's height.

Approach 2: Iterative In-order Traversal (High-ROI Interview Pattern)

While recursion is intuitive, an iterative solution using a stack is a fundamental pattern for tree problems in interviews. It avoids recursion depth limits and demonstrates a deeper command of traversal mechanics. The logic remains the same: simulate an in-order traversal and stop after visiting k elements.

This video from NeetCode provides an excellent step-by-step visualization and implementation of the iterative approach.

Here's the Python code, which is a template you should become very familiar with:

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

class Solution:
    def kthSmallest(self, root: TreeNode, k: int) -> int:
        stack = []
        current = root

        while current or stack:
            # 1. Go left and push nodes onto the stack
            while current:
                stack.append(current)
                current = current.left
            
            # 2. Pop the last node (the next smallest)
            current = stack.pop()
            
            # 3. "Visit" the node
            k -= 1
            if k == 0:
                return current.val
            
            # 4. Move to the right subtree
            current = current.right

• Time Complexity: In the worst case, we might need to visit k nodes, so the complexity is $O(H+k)$, where $H$ is the height of the tree (for the initial descent). If k is close to N, this is $O(N)$.
• Space Complexity: $O(H)$ for the stack.

      5
     / \
    3   7
   / \   \
  2   4   8

Finding the k-th Largest Element

Now, how can we adapt this to find the k-th largest element?

There are two excellent ways to solve this.

Method 1: Using the Total Node Count

The k-th largest element is equivalent to the $(N-k+1)$-th smallest element, where N is the total number of nodes in the tree. This approach would involve:

1. A first pass to count all nodes (N).
2. A second pass using one of our kthSmallest methods with the new k value of N - k + 1.

This is a valid solution, but it requires two passes over the tree.

Method 2: Reverse In-order Traversal

A more elegant, single-pass solution exists. If a regular in-order traversal (Left -> Node -> Right) gives us elements in ascending order, what would a reverse in-order traversal (Right -> Node -> Left) give us? Descending order!

Therefore, the k-th largest element is simply the k-th node we visit during a reverse in-order traversal.

The following resource explains this approach beautifully.

You can implement this both recursively and iteratively, just by swapping the left and right traversal logic.

Follow-up: Handling Frequent Queries

In a FAANG-level interview, a common follow-up question is: "What if the tree is modified frequently (many insertions/deletions) and you need to answer kthSmallest queries very fast?"

Our current approaches take $O(k)$ or $O(H+k)$ time for each query. If we have many queries, this becomes inefficient. The optimal solution involves augmenting the TreeNode structure.

Augmented Binary Search Tree:
We can modify the TreeNode class to include a field, say left_subtree_size, which stores the number of nodes in its left subtree.

class AugmentedTreeNode:
    def __init__(self, val):
        self.val = val
        self.left = None
        self.right = None
        self.left_subtree_size = 0 # or total subtree size

When we search for the k-th smallest element, we can use these counts to decide whether to go left, right, or stop:

1. Start at the root. Let the number of nodes in the left subtree be L (current.left_subtree_size).
2. The rank of the current node is L + 1.
3. If k == L + 1, the current node is our answer.
4. If k < L + 1, the element is in the left subtree. Recurse left, looking for the k-th smallest.
5. If k > L + 1, the element is in the right subtree. Recurse right, but now we're looking for the (k - (L + 1))-th smallest element in that subtree.

This augmented approach reduces the query time to $O(H)$, the height of the tree. The trade-off is that insertions and deletions become slightly more complex, as you need to update the subtree size counts all the way up to the root, but this is still an $O(H)$ operation.

Conclusion

Today we leveraged a core BST property to solve a classic interview problem.

Key Takeaways:

• The k-th smallest element in a BST is the k-th element in an in-order traversal (Left -> Node -> Right).
• The k-th largest element is the k-th element in a reverse in-order traversal (Right -> Node -> Left).
• For interviews, be prepared to implement these traversals both recursively and iteratively (using a stack). The iterative approach is a crucial, high-ROI pattern.
• For follow-up questions about frequent queries, the optimal solution is an augmented BST that stores subtree sizes, reducing query time from $O(N)$ to $O(H)$.

In our next lesson, we will tackle another BST construction problem: constructing a height-balanced BST from a sorted array. You will see how choosing the right element for the root at each step is key, a concept related to what we explored today.