Topics › Binary Search Trees › Practice Problems

Kth Smallest and Largest element in BST

Medium TUF ID 105 A2ZSDE

Problem

Given the root of a binary search tree with n distinct keys and an integer k, return {kthSmallest, kthLargest}: the key that would be at position k if all keys were sorted in increasing order, and the key at position k in decreasing order. If k is not between 1 and n, return {-1, -1}. (LeetCode asks only for the k-th smallest and guarantees 1 <= k <= n.)

Constraints: up to 10^4 nodes; keys fit in an int.

Example 1

            5
           / \
          3   6
         / \
        2   4
       /
      1

k = 3   answer = {3, 4}

The keys in order are 1 2 3 4 5 6. The third from the start is 3; the third from the end is 4.

Example 2

        3
       / \
      1   4          k = 1   answer = {1, 4}
       \
        2

With k = 1 the answers are simply the minimum and the maximum.

Approach 1: Brute force, collect and sort

Intuition. Ignore the tree structure: gather all keys in any order, sort them, and read off positions k - 1 and n - k (0-based). The k-th largest is the (n - k + 1)-th smallest, which sits at index n - k. Correct for any binary tree, but it pays for a sort the BST has already done for us.

Algorithm.

  1. Collect every key with a preorder traversal.
  2. Sort the list and let n be its size.
  3. If k < 1 or k > n, return {-1, -1}.
  4. Return {vals[k - 1], vals[n - k]}.
static int[] bruteForce(TreeNode root, int k) {
    List<Integer> vals = new ArrayList<>();
    preorder(root, vals);
    Collections.sort(vals);
    int n = vals.size();
    if (k < 1 || k > n) return new int[]{-1, -1};
    return new int[]{vals.get(k - 1), vals.get(n - k)};
}

static void preorder(TreeNode node, List<Integer> out) {
    if (node == null) return;
    out.add(node.val);
    preorder(node.left, out);
    preorder(node.right, out);
}

Complexity. Time O(n log n) for the sort. Space O(n) for the list.

Approach 2: Better, inorder is already sorted

Intuition. An inorder traversal (left subtree, node, right subtree) of a BST visits the keys in increasing order: every key in the left subtree is smaller than the node and every key in the right subtree is larger, and the same holds recursively. So the inorder list is the sorted list, with no sorting step.

Algorithm.

  1. Collect the keys with an inorder traversal.
  2. If k is out of range, return {-1, -1}.
  3. Return {sorted[k - 1], sorted[n - k]}.
static int[] better(TreeNode root, int k) {
    List<Integer> sorted = new ArrayList<>();
    inorder(root, sorted);
    int n = sorted.size();
    if (k < 1 || k > n) return new int[]{-1, -1};
    return new int[]{sorted.get(k - 1), sorted.get(n - k)};
}

static void inorder(TreeNode node, List<Integer> out) {
    if (node == null) return;
    inorder(node.left, out);
    out.add(node.val);
    inorder(node.right, out);
}

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

This still visits all n nodes and stores all keys, even when k = 1 and the answer is a few steps from the root.

Approach 3: Optimal, stop the traversal at the k-th node

Intuition. Inorder traversal produces keys one at a time in increasing order, so we only need to count: the k-th key produced is the answer, and the traversal can stop right there. To stop early cleanly, write the traversal iteratively with an explicit stack.

The iterative inorder works like this. From the current node, push it and every node down its left edge; the last one pushed has no left child, so it is the smallest key not yet visited. Pop it, count it, then move to its right child and repeat the “push the left edge” step there. The stack always holds the ancestors whose own key and right subtree are still pending.

For the k-th largest, run the mirror image, a reverse inorder (right subtree, node, left subtree), which produces keys in decreasing order: push the right edge, pop, count, move to the left child. This avoids having to know n, which would cost a full pass to compute.

Algorithm (k-th smallest).

  1. cur = root, count = 0, 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; increment count; if count == k, return its key;
    • move cur to the popped node’s right child.
  3. If the loop ends, there are fewer than k keys: return -1.

The k-th largest is identical with left and right swapped.

static int[] optimal(TreeNode root, int k) {
    return new int[]{kthSmallest(root, k), kthLargest(root, k)};
}

static int kthSmallest(TreeNode root, int k) {
    Deque<TreeNode> stack = new ArrayDeque<>();
    TreeNode cur = root;
    int count = 0;
    while (cur != null || !stack.isEmpty()) {
        while (cur != null) {                        // go as far left as possible
            stack.push(cur);
            cur = cur.left;
        }
        cur = stack.pop();                           // the next key in increasing order
        if (++count == k) return cur.val;
        cur = cur.right;
    }
    return -1;                                       // fewer than k keys (or k < 1)
}

static int kthLargest(TreeNode root, int k) {
    Deque<TreeNode> stack = new ArrayDeque<>();
    TreeNode cur = root;
    int count = 0;
    while (cur != null || !stack.isEmpty()) {
        while (cur != null) {                        // mirror image: go as far right as possible
            stack.push(cur);
            cur = cur.right;
        }
        cur = stack.pop();                           // the next key in decreasing order
        if (++count == k) return cur.val;
        cur = cur.left;
    }
    return -1;
}

Dry run of kthSmallest on Example 1 with k = 3:

step action stack (top at right) popped count
1 push left edge from 5 5, 3, 2, 1   0
2 pop; 1 has no right child 5, 3, 2 1 1
3 pop; 2 has no right child 5, 3 2 2
4 pop 5 3 3, return 3

And kthLargest with k = 3: push the right edge 5, 6; pop 6 (count 1); pop 5 (count 2) and move to its left child 3; push the right edge from 3, which is 3, 4; pop 4 (count 3), return 4.

Complexity. Time O(h + k): reaching the first key costs one walk down an edge of length at most h, and each later step pops one node and pushes nodes that are popped soon after. Space O(h) for the stack. For small k in a balanced tree this is O(log n), far less than O(n).

Two further refinements are worth knowing. Morris traversal threads temporary links through the tree to do inorder in O(1) extra space. And if the tree is modified often while k-th queries keep arriving (the LeetCode follow-up), store in each node the size of its subtree; then the k-th smallest is found by one walk from the root, comparing k with the left subtree’s size, in O(h) per query.

Edge cases and pitfalls