Dynamic Programming · Hard

321. Count Balanced Binary Trees of Height h

Topics: Dynamic Programming

Solve it on DeVert Campus - free code runner

Problem

A binary tree is height-balanced if, for every node, the heights of its left and right subtrees differ by at most 1. Given h, count the number of distinct height-balanced binary trees (structurally, ignoring node values) whose height is exactly h. Since this count grows doubly-exponentially, print it modulo 10^9 + 7.

Input and output

Input: A single line contains the integer h, the target height. Output: Print a single integer - the count of height-balanced trees of height exactly h, modulo 10^9 + 7.

Example

Input:
2
Output:
3

Constraints

0 <= h <= 1000

Key observation

Applying the modulo operation AFTER EVERY multiplication and addition (not just at the very end) is essential - waiting to apply modulo only at the end would already have overflowed the integer type long before that final step.

Intuition

Count the number of distinct height-balanced binary trees of exactly a given height. Brute force: Try to explicitly construct and enumerate every possible balanced tree shape directly - since this count grows DOUBLY-exponentially, actual enumeration becomes completely impossible even for modest h. Optimised: The recurrence directly encodes the THREE valid ways to combine two balanced subtrees into a height-i tree (both height i-1, or one height i-1 and one height i-2 in either arrangement) - building this bottom-up computes dp[h] in linear time WITHOUT ever constructing an actual tree, since only the COUNT is needed.

Bottom-up iterative DP with two rolling values

Build the sequence iteratively from the base cases upward: starting with prev2 = f(0) = 1 and prev1 = f(1) = 1, repeatedly compute the next value as prev1 * (2*prev2 + prev1) mod 10^9+7, then shift prev2 and prev1 forward, continuing until reaching h. Since each new value only ever depends on the previous two, there's no need to keep a full array of all computed values - just two rolling variables suffice, giving a single loop with constant extra memory. Time: O(h) Space: O(1)

Top-down memoized recursion

Use the exact same recursive relation f(h) = f(h-1) * (2*f(h-2) + f(h-1)), but before computing f(k) for any k, check a memo array/cache to see if it was already computed; if so return the cached value immediately, otherwise compute it recursively and store the result before returning. Since each distinct value of k from 0 to h is now computed exactly once, the exponential blowup of the naive version disappears, at the cost of an array to hold cached results and applying the modulo 10^9+7 at each multiplication to keep numbers bounded. Time: O(h) Space: O(h)

Naive recursion on the recurrence

Define f(h) directly from the structural recurrence f(h) = f(h-1) * (2*f(h-2) + f(h-1)), with base cases f(0) = f(1) = 1, and evaluate f(h) using plain recursive calls with no caching. Because f(h-1) and f(h-2) each independently re-derive their own f(h-2), f(h-3), and so on, the same smaller subproblems get recomputed an enormous number of times, similar to naive recursive Fibonacci but branching more aggressively. This is correct but becomes computationally infeasible well before h reaches even a few dozen. Time: Exponential Space: O(h) recursion depth

Step by step

1. dp[i] = number of height-balanced trees of height exactly i; dp[0]=1 (an empty tree, or single node, counts as height 0 here), dp[1]=1. 2. A tree of height i is formed by combining two balanced subtrees whose heights are (i-1,i-1), (i-1,i-2), or (i-2,i-1) - working out to dp[i] = dp[i-1] * (2*dp[i-2] + dp[i-1]). 3. Take every multiplication and addition MODULO 10^9+7 as you go, since the true values overflow ordinary integer types almost immediately (this count grows doubly-exponentially).

Complexity

Time: O(h) - building the sequence bottom-up from dp[0] and dp[1], with modular arithmetic at each step. Space: O(1) if only the previous two values are kept, or O(h) if the full sequence is stored.

Edge cases

Test h = 0 and h = 1 (the given base cases) and h = 1000 (upper bound), where every multiplication/addition must be taken modulo 10^9 + 7 to avoid overflow.

Interview tip

Be ready to explain WHY modular arithmetic is non-negotiable here, not just a formality - the true values overflow standard 64-bit integers extremely quickly, given the doubly-exponential growth rate explicitly mentioned in the problem.

All DSA problems