Find Kth smallest element in a BST

The problem, simply

You get a binary search tree (BST) and a number k, counting from 1. Give back the k-th smallest number in the tree. For k = 1 that is the smallest, for k = 2 the second smallest, and so on.

A story to remember it

Picture it

Examples

Way 1: the simple way

Write the whole sorted order down and pick position k - 1.

1. Visit left side, me, right side for every spot and add each number to a list.
2. Return list.get(k - 1).

import java.util.*;

class Solution {
    public int kthSmallest(TreeNode root, int k) {
        List<Integer> sorted = new ArrayList<>();
        inorder(root, sorted);
        return sorted.get(k - 1);
    }

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

How long it takes: every spot, O(n), even when k is 1. Extra room: the list, O(n).

Way 2: the clever way

Walk in sorted order with your own pile (a stack) and stop at the k-th spot.

1. Push the left road from the top: the top, its left child, that one's left child, and so on.
2. Pop the top of the pile. That is the next smallest number. Count it.
3. If the count is k, return it.
4. Otherwise push the left road starting from the popped spot's right child, and go back to step 2.

import java.util.*;

class Solution {
    public int kthSmallest(TreeNode root, int k) {
        Deque<TreeNode> pile = new ArrayDeque<>();
        TreeNode curr = root;
        while (curr != null || !pile.isEmpty()) {
            while (curr != null) {          // push the left road
                pile.push(curr);
                curr = curr.left;
            }
            curr = pile.pop();              // the next smallest
            if (--k == 0) return curr.val;  // that was the k-th
            curr = curr.right;              // its right side comes next
        }
        return -1;                          // k is bigger than the tree
    }
}

How long it takes: walking down the first left road is h steps, and then we pop k times. So O(h + k). Extra room: the pile is at most one road tall, O(h).

Another way: count sizes for many questions

If the tree changes often and you are asked for the k-th smallest many times, store in every spot how many spots its left side has (leftSize). Then one walk answers a question:

• If k == leftSize + 1, this spot is the answer.
• If k <= leftSize, go left.
• Otherwise go right with k - leftSize - 1.

class SizedNode {
    int val, leftSize;                       // how many spots are on my left side
    SizedNode left, right;
    SizedNode(int v) { val = v; }
}

class Solution {
    // insert keeps leftSize correct: going left means one more on that side
    SizedNode insert(SizedNode node, int v) {
        if (node == null) return new SizedNode(v);
        if (v < node.val) { node.leftSize++; node.left = insert(node.left, v); }
        else node.right = insert(node.right, v);
        return node;
    }

    int kthSmallest(SizedNode node, int k) {
        while (node != null) {
            if (k == node.leftSize + 1) return node.val;
            if (k <= node.leftSize) node = node.left;
            else { k -= node.leftSize + 1; node = node.right; }
        }
        return -1;
    }
}

Each question and each insert is O(h). In a balanced tree that is O(log n), even with millions of questions.

Watch out for

• Recursive versions need a shared counter (a field or an array). A plain int parameter is copied, so counts from the left side are lost.
• Recursive versions should stop early once the answer is found, or they quietly walk the whole tree.
• k larger than the tree size. Say what you return (here -1), or throw.
• A Morris walk gives O(1) room, but if you return early in the middle, threads are left in the tree. Finish the walk or undo them first.

What interviewers ask next

• The tree changes often and k-th smallest is asked often. Use the leftSize idea above. This is called an order statistic tree.
• k-th largest instead. Walk right side, me, left side. See "Find Kth largest element in a BST".
• Median of the BST. Count the spots, then find the middle one or two with this walk. See "Find the median of BST in O(n) time and O(1) space".
• k smallest numbers in a plain unsorted array. Use a max heap of size k, O(n log k), or quickselect, O(n) on average.

Remember this

Grown-up words: reading left, me, right is an inorder traversal, the pile is the stack of an iterative inorder walk, O(h + k) is "the height to get started plus k pops", and the spot that knows its left side's size is an augmented node in an order statistic tree.