Find the median of BST in O(n) time and O(1) space

The problem, simply

You get a binary search tree (BST) with n numbers. Give back the median: the middle number when they are sorted. If n is even there are two middle numbers, and the median is their average.

The hard part is the rule: use only O(1) extra room. No list of numbers, no stack, and no recursion (recursion secretly uses a stack as tall as the tree).

A story to remember it

Picture it

Examples

Way 1: the simple way

Write the sorted order into a list and look at the middle.

1. Visit left part, room, right part, adding each number to a list.
2. If n is odd, return list[n / 2]. If even, return the average of list[n/2 - 1] and list[n/2].

import java.util.*;

class Solution {
    public double findMedian(TreeNode root) {
        List<Integer> sorted = new ArrayList<>();
        inorder(root, sorted);
        int n = sorted.size();
        if (n == 0) return 0;
        if (n % 2 == 1) return sorted.get(n / 2);
        return (sorted.get(n / 2 - 1) + sorted.get(n / 2)) / 2.0;
    }

    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 once, O(n). Extra room: the list, O(n), plus O(h) for the recursion. This breaks the O(1) rule.

Way 2: the clever way

Walk in sorted order with ropes instead of a stack. This is the Morris walk. At each spot curr:

1. If curr has no left part, visit it and go right. (Going right may follow a rope back up.)
2. Otherwise find the last spot of its left part: go left once, then right until the next step is null or is curr.
3. If that spot's right is null, tie a rope: point it at curr, then go left.
4. If that spot's right is already curr, you came back by the rope. Untie it, visit curr, and go right.

Do the walk twice. The first time, only count, to learn n. The second time, write down the number at position (n + 1) / 2 and, for even n, also at n / 2 + 1. Let the second walk finish so every rope gets untied.

class Solution {
    public double findMedian(TreeNode root) {
        if (root == null) return 0;
        int n = walk(root, 0, 0, new int[2]);     // first walk: just count
        int a = (n + 1) / 2, b = n / 2 + 1;       // middle positions, from 1 (same if n is odd)
        int[] got = new int[2];
        walk(root, a, b, got);                    // second walk: pick them up
        return (got[0] + got[1]) / 2.0;
    }

    // Morris walk in sorted order. Saves the a-th and b-th numbers. Returns the count.
    private int walk(TreeNode root, int a, int b, int[] got) {
        int count = 0;
        TreeNode curr = root;
        while (curr != null) {
            if (curr.left != null) {
                TreeNode last = curr.left;                       // last spot of the left part
                while (last.right != null && last.right != curr) last = last.right;
                if (last.right == null) {                        // first time: tie a rope
                    last.right = curr;
                    curr = curr.left;
                    continue;
                }
                last.right = null;                               // back by the rope: untie
            }
            count++;                                             // visit curr
            if (count == a) got[0] = curr.val;
            if (count == b) got[1] = curr.val;
            curr = curr.right;
        }
        return count;
    }
}

How long it takes: each walk is O(n). Finding "the last spot of the left part" looks slow, but each branch is walked at most a few times in total across the whole walk. Two walks, still O(n). Extra room: a few fingers, O(1).

Watch out for

• Integer division for even n: (a + b) / 2 drops the half. Divide by 2.0, and be careful that a + b can overflow for huge values (use long).
• Off by one positions. With counting from 1, the middle is (n + 1) / 2 for odd n, and n / 2 and n / 2 + 1 for even n.
• Recursion or a stack is not O(1) room. The interviewer is asking for Morris on purpose.

What interviewers ask next

• Why does Morris not lose any spot? The left part's last spot never has a right child of its own, so its empty right arrow is free to borrow, and it is given back on the way up.
• Can you find the median in one walk? Only if you already know n, for example if the tree stores its size. Otherwise you need the count first.
• Median of a stream of numbers (LeetCode 295). Keep two heaps: a max heap for the smaller half and a min heap for the bigger half. O(log n) per number.
• The k-th smallest in O(1) room. Same Morris walk, stop counting at k, and still finish the walk to untie ropes.

Remember this

Grown-up words: exploring left, room, right is an inorder traversal, the rope is a thread (a temporary right pointer to the inorder successor), the whole trick is a Morris traversal, and the last spot of the left part is the inorder predecessor.