Kth Smallest Element in a BST

medium bst in-order stack

Problem

Given the root of a binary search tree, and an integer k, return the kth smallest element in the tree.

An in-order walk of a BST yields values in sorted order. Run an iterative in-order traversal with an explicit stack and stop the moment you have popped k nodes — no need to visit the rest.

Inputlevel-order: [5, 3, 6, 2, 4, _, _, 1], k = 3

Output3

In-order: 1, 2, 3, 4, 5, 6 — third value is 3.

level-order ('_' for null) — must be a BST

k (1-based)

def kth_smallest(root, k):
    stk = []
    cur = root
    while stk or cur:
        while cur:
            stk.append(cur); cur = cur.left
        cur = stk.pop()
        k -= 1
        if k == 0: return cur.val
        cur = cur.right

function kthSmallest(root, k) {
  const stk = []; let cur = root;
  while (stk.length || cur) {
    while (cur) { stk.push(cur); cur = cur.left; }
    cur = stk.pop();
    if (--k === 0) return cur.val;
    cur = cur.right;
  }
}

class Solution {
    public int kthSmallest(TreeNode root, int k) {
        Deque<TreeNode> stk = new ArrayDeque<>();
        TreeNode cur = root;
        while (!stk.isEmpty() || cur != null) {
            while (cur != null) { stk.push(cur); cur = cur.left; }
            cur = stk.pop();
            if (--k == 0) return cur.val;
            cur = cur.right;
        }
        return -1;
    }
}

int kthSmallest(TreeNode* root, int k) {
    stack<TreeNode*> stk;
    TreeNode* cur = root;
    while (!stk.empty() || cur) {
        while (cur) { stk.push(cur); cur = cur->left; }
        cur = stk.top(); stk.pop();
        if (--k == 0) return cur->val;
        cur = cur->right;
    }
    return -1;
}

Explanation

The key fact is that an inorder traversal (left, node, right) of a BST visits values in sorted order. So the kth value we visit is exactly the kth smallest, and we can stop as soon as we have counted k of them.

Instead of recursion, this solution does the inorder walk with an explicit stack. From the current node we keep pushing and going left until there is no left child, so the stack holds the chain down to the smallest unvisited value.

Then we pop a node — that is the next value in sorted order — and decrement k. If k hits 0, this popped node's value is the answer. Otherwise we move to its right child and repeat the push-left process.

Example: tree [5,3,6,2,4,_,_,1], k = 3. We push 5,3,2,1, then pop 1 (k=2), pop 2 (k=1), pop 3 (k=0) and return 3.

We only descend the left spine once plus pop k times, so the work is about O(h + k) with O(h) stack space, and we never traverse the whole tree when k is small.

Time: O(h + k) Space: O(h) stack