Topics › Binary Search Trees › Practice Problems

Find K-th smallest element in BST

Medium TUF ID 2776 SDEBlind 75

Problem

Given the root of a binary search tree and an integer k, return the k-th smallest key in the tree, counting from 1. So k = 1 asks for the minimum and k = n for the maximum. The code below returns -1 when the tree has fewer than k nodes; LeetCode guarantees that never happens.

Constraints (LeetCode 230): 1 ≤ k ≤ n ≤ 10^4, keys between 0 and 10^4, all keys distinct.

Example 1

      3
     / \
    1   4
     \
      2

k = 1   answer = 1

The sorted keys are 1, 2, 3, 4; the first is 1.

Example 2

          5
         / \
        3   6
       / \
      2   4
     /
    1

k = 3   answer = 3

The sorted keys are 1, 2, 3, 4, 5, 6; the third is 3.

Approach 1: Brute force, collect and sort

Intuition. Forget that the tree is a BST. Any traversal gives us all the keys; sorting them and taking index k - 1 answers the question. This is the version you would write for an arbitrary binary tree, and it is a useful baseline because it is obviously correct.

Algorithm.

  1. Traverse the tree in any order and put every key into a list.
  2. Sort the list.
  3. Return list[k - 1], or -1 if k is out of range.
static int bruteForce(TreeNode root, int k) {
    List<Integer> keys = new ArrayList<>();
    collect(root, keys);
    Collections.sort(keys);
    return k >= 1 && k <= keys.size() ? keys.get(k - 1) : -1;
}

static void collect(TreeNode node, List<Integer> out) {
    if (node == null) return;
    out.add(node.val);                       // any traversal order works: we sort afterwards
    collect(node.left, out);
    collect(node.right, out);
}

Complexity. Time O(n log n) for the sort. Space O(n) for the list, plus O(h) recursion stack.

The sort is wasted work: the BST already stores its keys in order. We just need to read them in the right sequence.

Approach 2: Better, recursive in-order with an early stop

Intuition. An in-order traversal (left subtree, node, right subtree) of a BST visits keys in ascending order. That follows directly from the BST property: everything in the left subtree is smaller than the node and everything in the right subtree is larger, so the node must be visited exactly between them. Therefore the k-th node visited in-order is the k-th smallest key.

We do not need the whole sequence. Keep a countdown starting at k; decrement it at each visit; when it reaches zero, record the key and stop exploring.

Algorithm.

  1. Keep state = {remaining, answer} with remaining = k, answer = -1. An array is used so the recursive calls share one mutable counter.
  2. visit(node): if node is null or remaining is already 0, return.
  3. Recurse left. If that found the answer, return.
  4. Decrement remaining. If it hits 0, store node.val and return.
  5. Recurse right.
static int recursiveInorder(TreeNode root, int k) {
    int[] state = {k, -1};                   // {visits still needed, answer}
    visit(root, state);
    return state[1];
}

static void visit(TreeNode node, int[] state) {
    if (node == null || state[0] == 0) return;
    visit(node.left, state);
    if (state[0] == 0) return;               // the answer was found in the left subtree
    if (--state[0] == 0) {                   // this node is the k-th in sorted order
        state[1] = node.val;
        return;
    }
    visit(node.right, state);
}

Complexity. Time O(h + k): we descend at most h levels to the minimum and then visit k nodes, with each extra step along the way bounded by those. Space O(h) for the recursion.

When k is small and the tree is balanced, this touches only a sliver of the tree. The cost is the call stack, which is O(n) for a chain-shaped tree.

Approach 3: Optimal, iterative in-order with a stack

Intuition. The same in-order traversal can be driven by an explicit stack, which makes “stop at the k-th node” a plain return instead of a flag threaded through recursion.

The stack holds the nodes whose left subtree we are still processing, that is, the ancestors we have not visited yet. To find the next smallest key, walk left from the current node as far as possible, pushing every node on the way. The top of the stack is then the smallest unvisited key. Pop it, count it, and continue from its right child, because everything in that right subtree comes next in sorted order.

Algorithm.

  1. Set cur = root and create an empty stack.
  2. While cur is not null or the stack is not empty:
    • Push cur and move left until cur is null.
    • Pop a node. Decrement k; if it is now 0, return the node’s key.
    • Set cur to the popped node’s right child.
  3. If the loop ends, the tree had fewer than k nodes: return -1.
static int iterativeInorder(TreeNode root, int k) {
    Deque<TreeNode> stack = new ArrayDeque<>();
    TreeNode cur = root;
    while (cur != null || !stack.isEmpty()) {
        while (cur != null) {                // walk to the smallest key not yet visited
            stack.push(cur);
            cur = cur.left;
        }
        TreeNode node = stack.pop();
        if (--k == 0) return node.val;       // pops come out in ascending order
        cur = node.right;
    }
    return -1;
}

Dry run on Example 2 with k = 3 (stack written bottom to top):

step stack popped k after
walk left from 5 5, 3, 2, 1   3
pop 5, 3, 2 1 2
1.right is null, pop 5, 3 2 1
2.right is null, pop 5 3 0, return 3

Complexity. Time O(h + k): the first walk down is at most h pushes, and every popped node costs O(1) amortised because each node is pushed and popped at most once. Space O(h) for the stack.

This is the answer interviewers expect for LeetCode 230. It stops as soon as the answer is known, and its memory is bounded by the height rather than the size of the tree.

Approach 4: Morris traversal, constant extra space

Intuition. The stack in Approach 3 exists only to remember “where do I go after finishing this left subtree?”. Morris traversal stores that information inside the tree itself. Before descending into the left subtree of cur, find cur’s in-order predecessor (the rightmost node of the left subtree). Its right pointer is null, so we temporarily point it at cur. When the traversal later finishes the left subtree, it follows that thread back to cur. Finding the thread already present tells us the left side is done, so we remove it, visit cur, and go right.

There is a catch. If we stopped at the k-th node, threads created higher up would still be in place and the caller’s tree would contain cycles. So this version records the answer and keeps going until the traversal ends, which removes every thread.

Algorithm.

  1. cur = root, count = 0.
  2. While cur is not null:
    • If cur has no left child: visit it (increment count, record if count == k) and go right.
    • Otherwise find the predecessor pred. If pred.right is null, set it to cur and go left. If it already points to cur, reset it to null, visit cur, and go right.
  3. Return the recorded answer.
static int morris(TreeNode root, int k) {
    int count = 0, answer = -1;
    TreeNode cur = root;
    while (cur != null) {
        if (cur.left == null) {
            if (++count == k) answer = cur.val;
            cur = cur.right;                 // a real right child or a thread back to the successor
        } else {
            TreeNode pred = cur.left;        // in-order predecessor: rightmost node of the left subtree
            while (pred.right != null && pred.right != cur) pred = pred.right;
            if (pred.right == null) {
                pred.right = cur;            // first arrival: leave a thread so we can come back
                cur = cur.left;
            } else {
                pred.right = null;           // second arrival: left subtree done, remove the thread
                if (++count == k) answer = cur.val;
                cur = cur.right;
            }
        }
    }
    return answer;
}

Complexity. Time O(n): every edge is walked a constant number of times (the predecessor searches add up to O(n) overall). Space O(1) extra. The Java test checks that the tree’s shape is identical before and after the call.

Use this when memory is the binding constraint. Otherwise Approach 3 is better: it can stop early and never mutates shared data.

Follow-up from LeetCode. If the tree is modified often and k-th queries are frequent, store in every node the size of its left subtree. At a node with leftSize keys on its left: if k == leftSize + 1 the node is the answer; if k <= leftSize go left; otherwise subtract leftSize + 1 from k and go right. Each query becomes O(h), and inserts and deletes update the counts along their path.

Edge cases and pitfalls