Kth Smallest Element in a BST

▶

Steps

setup

Initialize stack and current pointer

We start with an empty stack and set current to the root node (value 3). The count of visited nodes is zero.

💡 Initialization sets up the data structures needed for iterative in-order traversal.

Line:

stack = []
current = root
count = 0

💡 The traversal will begin from the root, and the stack will help us backtrack after reaching leftmost nodes.

fill_cells

Push current node (3) onto stack and move left

Current node 3 is pushed onto the stack. We move current pointer to its left child (node 1).

💡 Pushing nodes onto the stack delays visiting them until their left subtree is fully explored.

Line:

while current:
    stack.append(current)
    current = current.left

💡 Stack keeps track of nodes to visit after left descendants are processed.

fill_cells

Push current node (1) onto stack and move left

Current node 1 is pushed onto the stack. We move current pointer to its left child, which is null.

💡 When current becomes null, it means we reached the leftmost node of this path.

Line:

while current:
    stack.append(current)
    current = current.left

💡 Reaching null signals no more left children; time to visit nodes from stack.

traverse

Pop node (1) from stack and visit it

Current is null, so we pop node 1 from the stack to visit it. Count increments to 1.

💡 Visiting a node means processing it after all left descendants are done.

Line:

current = stack.pop()
count += 1

💡 The first visited node in in-order traversal is the smallest node in the BST.

compare

Check if count equals k

Count is 1, which equals k=1. The algorithm returns the current node's value 1 as the kth smallest element.

💡 The stopping condition ensures we return the correct kth smallest value immediately.

Line:

if count == k:
    return current.val

💡 The algorithm stops as soon as the kth smallest node is visited, optimizing performance.

reconstruct

Return the kth smallest element

The algorithm returns 1 as the kth smallest element and terminates.

💡 Returning the value completes the search and provides the final answer.

Line:

return current.val

💡 The traversal does not continue after finding the kth smallest, saving time.

prune

Additional step: Explanation of why traversal stops early

Since k=1, the algorithm stops after visiting the first node in in-order traversal, which is the smallest element.

💡 Early stopping is a key optimization in kth smallest element search.

Line:

if count == k:
    return current.val

💡 Traversal does not need to visit all nodes, saving time especially in large BSTs.

expand

Hypothetical continuation if k was larger: move to right child

If count was less than k, the algorithm would move current pointer to the right child of the popped node to continue traversal.

💡 Moving to right child explores nodes greater than the current node in in-order traversal.

Line:

current = current.right

💡 Traversal order is left → node → right, so right subtree is visited after node.

fill_cells

Hypothetical push of right child (4) onto stack

If traversal continued, current node 4 would be pushed onto the stack and traversal would move left to node 2.

💡 Pushing right child nodes onto stack continues the in-order traversal pattern.

Line:

while current:
    stack.append(current)
    current = current.left

💡 Left children of right subtree are visited before the right subtree root.

reconstruct

Summary: In-order traversal visits nodes in ascending order

The in-order traversal order for this BST is [1, 2, 3, 4]. The algorithm visits nodes in this order, counting until it reaches k.

💡 Understanding this order is key to grasping why in-order traversal finds the kth smallest element.

💡 In-order traversal of BST yields sorted node values, enabling direct access to kth smallest by counting visits.

from typing import Optional

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

def kthSmallest(root: Optional[TreeNode], k: int) -> int:
    stack = []  # STEP 1: Initialize empty stack
    current = root  # STEP 1: Start at root
    count = 0  # STEP 1: Initialize count

    while current or stack:  # STEP 2: Continue while nodes remain
        while current:  # STEP 3: Push all left children
            stack.append(current)  # STEP 3: Push current
            current = current.left  # STEP 3: Move left
        current = stack.pop()  # STEP 4: Pop from stack
        count += 1  # STEP 4: Increment count
        if count == k:  # STEP 5: Check if kth smallest
            return current.val  # STEP 5: Return value
        current = current.right  # STEP 6: Move to right child

# Driver code
if __name__ == '__main__':
    root = TreeNode(3)
    root.left = TreeNode(1)
    root.right = TreeNode(4)
    root.left.right = TreeNode(2)
    print(kthSmallest(root, 1))  # Output: 1

📊

Kth Smallest Element in a BST - Watch the Algorithm Execute, Step by Step

Watching each stack operation and pointer movement reveals how in-order traversal naturally visits nodes in ascending order, making it the fastest way to understand the kth smallest element search.

Current node: 1

3124

Current node: 2

3124

3124

Current node: 2

3124

Result: [1]

Current node: 2

3124

Result: [1]

Return: 1

3124

Result: [1]

Return: 1

3124

Result: [1]

Return: 1

Current node: 5

3124

Result: [1]

Return: 1

3124

Result: [1]

Return: 1

3124

Result: [1, 2, 3, 4]

Return: 1

Key Takeaways

✓ In-order traversal visits BST nodes in ascending order, enabling direct access to the kth smallest element by counting visits.

This insight is hard to see from code alone because the traversal order and counting are implicit; visualization makes it explicit.

✓ The stack stores nodes whose left subtrees are fully explored but whose right subtrees are yet to be visited, enabling backtracking.

Understanding the stack's role clarifies how iterative traversal mimics recursion.

✓ The algorithm stops immediately when the kth smallest node is visited, optimizing performance by avoiding unnecessary traversal.

This early stopping is subtle in code but critical for efficiency, and visualization highlights this decision clearly.