Kth Smallest Element In a Bst

The problem

Given root, the top node of a binary search tree, and a number k, return the kth smallest value in the tree. Counting starts at 1, so k = 1 asks for the smallest value. Each node has a number val and two children, left and right; a missing child is null. In a binary search tree, all values in a node's left subtree (its left child and everything below that child) are smaller than the node's value, and all values in its right subtree are larger. No value appears twice, and the tree has at least k nodes.

What to write

type TreeNode = {val:number;left:TreeNode|null;right:TreeNode|null};

function solve(root: TreeNode | null, k: number): number

Input
root = {"val":3,"left":{"val":1,"left":null,"right":{"val":2,"left":null,"right":null}},"right":{"val":4,"left":null,"right":null}}
k = 1

Output
1

Input
root = {"val":5,"left":{"val":3,"left":{"val":2,"left":{"val":1,"left":null,"right":null},"right":null},"right":{"val":4,"left":null,"right":null}},"right":{"val":6,"left":null,"right":null}}
k = 3

Output
3

Constraints

k is between one and the node count.

Hints

Each hint gives away a little more. Open one only when you need it.

1. Hint 1: Use the constraints

   A binary search tree is already organized: at every node you know which side holds the smaller values. So nothing has to be sorted, and for a small k you should not need to look at the whole tree.

2. Hint 2: Ask a smaller question

   Stand at the root of the second example. Where in the tree is the smallest value, and how do you walk there? Once you have it, which node holds the second smallest value?

3. Hint 3: Name the pattern

   You are asked for a position in sorted order, and a binary search tree keeps its values in order. Reading a node's left subtree, then the node, then its right subtree gives the values from smallest to largest. That is binary search tree ordering, and the walk is called an inorder traversal. Read the lesson on binary search tree ordering.

4. Hint 4: State the invariant

   A node is counted after everything in its left subtree and before everything in its right subtree. So the nodes are counted from the smallest value to the largest, and the kth node counted holds the answer.

5. Hint 5: Sketch the pseudocode

   Keep a stack and a current node, starting at the root. Repeat: while there is a current node, push it and move to its left child. Then pop a node and lower k by 1. If k is now 0, return that node's value. Otherwise continue with its right child.

Solution

From here on the page gives the method away. If you are still solving, go back to the hints.

Approach

Inorder traversal visits a BST in ascending order. Stop as soon as the kth node is visited.

The slow way

Visit every node, collect all the values in a list, sort the list and return the entry at position k minus 1. It is correct, but it reads the whole tree and then sorts values that the tree already keeps in order, even when k is 1.

The insight

A binary search tree is sorted if you read it in the right order: everything in a node's left subtree, then the node, then everything in its right subtree. Reading a tree in this order is called an inorder traversal. Read like that, the values come out from smallest to largest, so you can count them as they come and stop at the kth.

Step by step

1. Keep a stack of the nodes passed on the way down. On the way left you pass nodes whose turn comes later. A stack is a list where the last item put on is the first one taken off, so it hands them back nearest first.
2. Go left as far as you can, stacking each node you pass. Smaller values are always on the left, so the smallest value that remains is at the end of the leftmost path from here.
3. Take the top node: it holds the smallest value not yet counted. Its left subtree is empty or already counted. Everything else that remains is larger: its own right subtree, and the nodes further down the stack with their right subtrees.
4. Count it, and stop when it is the kth. Nodes leave the stack from the smallest value to the largest, so the first one off is the smallest and the kth one off is the answer.
5. Continue in its right subtree, where the next values are. The values in its right subtree are larger than this node but smaller than every node still on the stack, so they come next.

The same numbered steps are the comments in each solution.

function solve(root, k) {
  // 1. Keep a stack of the nodes passed on the way down
  const stack = [];
  let cur = root;
  while (cur || stack.length) {
    // 2. Go left as far as you can, stacking each node you pass
    while (cur) {
      stack.push(cur);
      cur = cur.left;
    }
    // 3. Take the top node: it holds the smallest value not yet counted
    cur = stack.pop();
    // 4. Count it, and stop when it is the kth
    if (--k === 0) return cur.val;
    // 5. Continue in its right subtree, where the next values are
    cur = cur.right;
  }
  // k is promised to be valid, so this line is never reached.
  throw Error("Invalid k");
}

type TreeNode = { val: number; left: TreeNode | null; right: TreeNode | null };
function solve(root: TreeNode | null, k: number): number {
  // 1. Keep a stack of the nodes passed on the way down
  const stack: TreeNode[] = [];
  let cur = root;
  while (cur || stack.length) {
    // 2. Go left as far as you can, stacking each node you pass
    while (cur) {
      stack.push(cur);
      cur = cur.left;
    }
    // 3. Take the top node: it holds the smallest value not yet counted
    cur = stack.pop()!;
    // 4. Count it, and stop when it is the kth
    if (--k === 0) return cur.val;
    // 5. Continue in its right subtree, where the next values are
    cur = cur.right;
  }
  // k is promised to be valid, so this line is never reached.
  throw Error("Invalid k");
}

def solve(root, k):
    # 1. Keep a stack of the nodes passed on the way down
    stack = []
    node = root
    while node or stack:
        # 2. Go left as far as you can, stacking each node you pass
        while node:
            stack.append(node)
            node = node.left
        # 3. Take the top node: it holds the smallest value not yet counted
        node = stack.pop()
        # 4. Count it, and stop when it is the kth
        k -= 1
        if k == 0:
            return node.val
        # 5. Continue in its right subtree, where the next values are
        node = node.right
    # k is promised to be valid, so this line is never reached.
    raise ValueError("Invalid k")

Example 2, traced

Input
root = {"val":5,"left":{"val":3,"left":{"val":2,"left":{"val":1,"left":null,"right":null},"right":null},"right":{"val":4,"left":null,"right":null}},"right":{"val":6,"left":null,"right":null}}
k = 3

Output
3

Why it works

A node leaves the stack only when nothing smaller remains: its left subtree has been counted, and its right subtree and every node still on the stack are larger. So the nodes are counted from smallest to largest with none skipped, and the kth one counted is the kth smallest.

Time and space

Time O(h + k), extra space O(h). A node is stacked at most once. When the loop stops, k nodes have been taken off and the rest lie on one path down the tree, which has at most h nodes, the tree's height. So the time is O(h + k) and the space is O(h).

n and m are input sizes, h is a tree's height, V and E are a graph's vertices and edges. The bounds describe this reference solution.

Watch out

• Count a node when you take it off the stack, not when you put it on. Counting on the way down counts 5, 3 and 2 in the second example and returns 2 instead of 3.
• After counting a node, continue with its right child. A version that forgets this still passes both examples, but on the first example's tree with k = 2 it returns 3 and never sees the 2.
• k counts from 1. If you put the values into a list in rising order, the answer is at position k minus 1. Position k holds the next larger value: 2 instead of 1 in the first example.

Pattern

Binary search tree ordering. Builds on Lowest Common Ancestor of a Binary Search Tree and Validate Binary Search Tree.

Keep going