Find Kth largest 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 largest number. For k = 1 that is the biggest number, for k = 2 the second biggest, and so on.

A story to remember it

Picture it

Examples

Way 1: the simple way

Write the sorted order down, smallest first, and count back from the end.

1. Visit left side, me, right side and add each number to a list.
2. The k-th largest is at position size - k.

import java.util.*;

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

    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). Extra room: the list, O(n).

Way 2: the clever way

Walk biggest first with a pile, and stop at the k-th pop. It is the kth smallest walk with left and right swapped.

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

import java.util.*;

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

How long it takes: h steps down the first right road, then k pops, O(h + k). Extra room: the pile, O(h).

Another way: recursion with a shared counter

The same walk written as recursion. The counter must live outside the calls so that all of them share it.

class Solution {
    private int left, answer;

    public int kthLargest(TreeNode root, int k) {
        left = k;
        answer = -1;
        walk(root);
        return answer;
    }

    private void walk(TreeNode node) {
        if (node == null || left == 0) return;   // stop early once found
        walk(node.right);
        if (left == 0) return;
        if (--left == 0) { answer = node.val; return; }
        walk(node.left);
    }
}

Same O(h + k) time and O(h) room for the call stack.

Watch out for

• Swapping only one of the two directions in Way 2. Push the right road, and after a pop move to the left child. Swapping just one gives a wrong order.
• A counter passed as a plain int parameter in recursion. Each call gets its own copy.
• k bigger than the tree, or k <= 0. Decide what to return.

What interviewers ask next

• Do it with O(1) room. Use a mirrored Morris walk: before going right from a spot, find the smallest spot on its right side (right once, then left as far as you can) and point that spot's empty left arrow back up, like a rope. Climb back by the rope, untie it, count the spot, then go left. Untie any ropes left behind if you stop early.
• Many questions on a changing tree. Store subtree sizes in each spot, then the k-th largest is the (n - k + 1)-th smallest, found in O(h).
• k-th largest in a stream of numbers (LeetCode 703). Keep a min heap of the k biggest seen so far. Its top is the answer.
• Sum of the k biggest numbers. Same walk, add instead of only counting.

Remember this

Grown-up words: right, me, left is a reverse inorder traversal, the pile is the stack of an iterative traversal, the shared counter is a field shared across recursive calls, and O(h + k) means the height to start plus k pops.