Count BST nodes that lie in a given range

The problem, simply

You get a binary search tree (BST) and two numbers, low and high. Count how many numbers in the tree are between them, including low and high themselves.

A story to remember it

Picture it

Examples

Way 1: the simple way

Visit every spot and check it.

1. If the spot is null, return 0.
2. Count 1 if low <= val <= high, else 0.
3. Add the counts from both sides.

class Solution {
    public int countInRange(TreeNode root, int low, int high) {
        if (root == null) return 0;
        int here = (root.val >= low && root.val <= high) ? 1 : 0;
        return here + countInRange(root.left, low, high)
                    + countInRange(root.right, low, high);
    }
}

How long it takes: every spot, O(n). Extra room: the recursion, O(h). Correct on any tree, but it ignores the order.

Way 2: the clever way

Use the order to skip whole sides.

1. If the spot is null, return 0.
2. If val < low, this spot and its whole left side are too small. Only look at the right side.
3. If val > high, this spot and its whole right side are too big. Only look at the left side.
4. Otherwise count this spot and look at both sides.

class Solution {
    public int countInRange(TreeNode root, int low, int high) {
        if (root == null) return 0;
        if (root.val < low) return countInRange(root.right, low, high);   // left side is too small
        if (root.val > high) return countInRange(root.left, low, high);   // right side is too big
        return 1 + countInRange(root.left, low, high)                     // I fit
                 + countInRange(root.right, low, high);
    }
}

How long it takes: we visit the k spots that fit, plus at most two roads of spots that do not (one along the low edge, one along the high edge). So O(h + k). Extra room: the recursion, O(h).

Watch out for

• Both ends count: use >= and <=.
• A spot that is too small is not counted, but its right side might still hold numbers that fit. Do not stop there.
• For the sum version (LeetCode 938, range sum of BST), add root.val instead of 1, and use long if the sum can be large.

What interviewers ask next

• Return the numbers, not the count, in sorted order. Same pruning, but walk left side, me, right side and add to a list.
• Many range questions on a fixed tree. Store in each spot the size of its part of the tree. Then "how many are <= high" is one road down, O(h), and the answer is countAtMost(high) - countAtMost(low - 1).
• Remove every spot outside the range (LeetCode 669, trim a BST). Too small: return the trimmed right side. Too big: return the trimmed left side. Otherwise trim both children.
• Count how many full subtrees lie entirely inside the range. Each spot returns whether its whole part fits. A part fits if both children's parts fit and the spot fits.

Remember this

Grown-up words: skipping whole sides is pruning, the low and high marks are the bounds of a range query, k is the size of the answer, and storing part sizes in each spot makes an order statistic tree.