Count Balanced Binary Trees of Height h

The problem, simply

A binary tree is balanced when, at every node, the height of the left side and the height of the right side differ by at most 1. Here height means the number of rows: a single node has height 1, and the empty tree has height 0.

You get h. Count how many different shapes of balanced binary tree have height exactly h. Only the shape matters, not the values. The count grows very fast, so give it modulo 10^9 + 7.

A story to remember it

Picture it

Examples

Way 1: the simple way

Turn the story into a rule. Let count(h) be the answer. The top node uses sides of height (h-1, h-1), (h-1, h-2) or (h-2, h-1). Each side can be any shape of its height, so:

count(h) = count(h-1) * count(h-1) + 2 * count(h-1) * count(h-2)

with count(0) = 1 (empty) and count(1) = 1 (one node).

class Solution {
    static final long MOD = 1_000_000_007L;

    long count(int h) {
        if (h <= 1) return 1;                    // empty tree, or a single node
        long a = count(h - 1);                   // shapes one row shorter
        long b = count(h - 2);                   // shapes two rows shorter
        // (a, a) + (a, b) + (b, a)
        return (a * a % MOD + 2 * a % MOD * b % MOD) % MOD;
    }

    public long countBT(int h) {
        return count(h);
    }
}

How long it takes: each call makes two calls, a lot like Fibonacci, so the calls grow up to O(2^h). Extra room: the chain of waiting calls, O(h).

Way 2: the clever way

The tree above asks for count(3) twice and count(2) three times. Fill an array from the bottom instead: each height only needs the two heights below it.

1. count[0] = 1, count[1] = 1.
2. For i from 2 to h: count[i] = count[i-1]^2 + 2 * count[i-1] * count[i-2], modulo 10^9 + 7.
3. Return count[h].

class Solution {
    public long countBT(int h) {
        final long MOD = 1_000_000_007L;
        if (h <= 1) return 1;
        long[] count = new long[h + 1];
        count[0] = 1;                            // the empty tree
        count[1] = 1;                            // one node
        for (int i = 2; i <= h; i++) {
            long a = count[i - 1], b = count[i - 2];
            long same = a * a % MOD;             // both sides one row shorter
            long mixed = 2 * a % MOD * b % MOD;  // one side two rows shorter, either side
            count[i] = (same + mixed) % MOD;
        }
        return count[h];
    }
}

How long it takes: one step per height, O(h). Extra room: the array, O(h). Only the last two values are used, so two variables give O(1) room.

Watch out for

• Height in rows vs height in edges. Here one node is height 1. If your judge says one node is height 0, everything shifts by one.
• count(0) = 1, not 0. The empty side is one valid choice, and it is what makes count(2) = 3.
• Do not forget the factor 2: (h-1, h-2) and (h-2, h-1) are mirror images and both count.

What interviewers ask next

• Fewest nodes in a balanced tree of height h? N(h) = N(h-1) + N(h-2) + 1, which is why balanced (AVL) trees have height about 1.44 * log2(n).
• Count balanced trees with exactly n nodes? You need both height and size: dp[h][n] sums over how many nodes go left. That is a 2D DP.
• Count all binary tree shapes with n nodes, no balance rule? That is the Catalan number C(n).
• Why is this called a balanced tree? It is the shape rule of an AVL tree, which keeps search fast.

Remember this

Grown-up words: a fair coat rack is a height-balanced binary tree (AVL shape), counting by multiplying the sides is the product rule, the repeated count(3) calls are overlapping subproblems, and the array is bottom-up DP.