7.43

From The Algorithm Design Manual Solution Wiki
Jump to navigation Jump to search

Idea: An in-order (left → root → right) traversal of a Binary Search Tree (BST) visits the keys in ascending order. Counting the visited nodes while traversing stops exactly at the i-th element, which is the required node.

Recursive Python code (1-based index i):
def kth_inorder(root, i):
    # returns the TreeNode that is the i-th smallest key
    counter = {"k": i} # mutable wrapper so recursion can update it
    
    def dfs(node):
        if not node:
            return None
        # 1. left subtree
        left = dfs(node.left)
        if left: # answer already found
            return left
        # 2. current node
        counter["k"] -= 1
        if counter["k"] == 0:
            return node
        # 3. right subtree
        return dfs(node.right)
    
    return dfs(root)

Iterative version (explicit stack):
def kth_inorder_iter(root, i):
    stack = []
    curr = root
    while stack or curr:
        # go to leftmost node
        while curr:
            stack.append(curr)
            curr = curr.left
        curr = stack.pop()
        i -= 1
        if i == 0:
            return curr
        curr = curr.right

Complexities: • Time – O(h + i) in worst case for skewed tree (traverses each visited node once); with a balanced tree O(log n + i). • Space – O(h) for the recursion/stack (h = tree height).

Optional optimization for many queries: store size = number of nodes in each subtree inside every node. Then: 1. Let left_size = node.left.size if node.left else 0. 2. If i == left_size + 1 → answer found. 3. If i ≤ left_size → recurse left. 4. Else recurse right with i − left_size − 1. This gives O(h) time per query with no extra traversal.


Back to Chapter 7