Count Ways to Build Rooms in an Ant Colony: Tree Combinatorics

You have an ant colony with n rooms numbered 0 to n-1. Room 0 is the starting room. Each other room has exactly one prerequisite room that must be built before it. Given prevRoom[i] (the prerequisite for room i, with prevRoom[0] = -1 for the root), count the number of different valid orders to build all the rooms. Return the answer modulo 10^9 + 7.

This is LeetCode 1916: Count Ways to Build Rooms in an Ant Colony, and it is one of those problems that looks like a graph traversal at first glance but turns into a combinatorics puzzle once you see the structure.

The problem

The prerequisite relationships form a rooted tree. Room 0 is the root, and every other room has exactly one parent. You need to find how many topological orderings of this tree exist. A topological ordering is any sequence where every room appears after its prerequisite.

For a linear chain like 0 -> 1 -> 2 -> 3, there is exactly one valid order. But when a node has multiple children, their subtrees can be interleaved in many ways, and counting those interleavings is the core challenge.

The approach

The key insight is that this problem reduces to counting the number of valid interleavings of independent subtree build sequences. At each node with children, you need to merge the build orders of each child's subtree while preserving the internal ordering within each subtree.

Multinomial coefficients for interleavings

Suppose a node has children with subtree sizes s1, s2, ..., sk. The total number of descendants to schedule is s1 + s2 + ... + sk. The number of ways to interleave k sequences of lengths s1, s2, ..., sk (while keeping each sequence in its original order) is the multinomial coefficient:

(s1 + s2 + ... + sk)! / (s1! * s2! * ... * sk!)

This is the same formula you would use to count distinct permutations of a multiset. Each child's subtree has a fixed internal order (because of its own dependencies), and you are choosing positions in the merged sequence for each group.

Modular arithmetic

Since the answer can be huge, you work modulo 10^9 + 7. Division modulo a prime is done via modular inverse: instead of dividing by x, you multiply by x^(p-2) mod p (Fermat's little theorem). You precompute all factorials and inverse factorials up to n so that every multinomial coefficient is a constant-time lookup.

DFS to combine everything

Run a post-order DFS. At each node:

1. Recursively compute the subtree size and the number of valid orderings for each child's subtree
2. Compute the multinomial coefficient for interleaving the children's subtrees
3. Multiply the multinomial coefficient by the product of ways for each child

The answer at the root is the total number of valid build orders.

def way_to_build_rooms(prev_room: list[int]) -> int:
    MOD = 10**9 + 7
    n = len(prev_room)

    children: list[list[int]] = [[] for _ in range(n)]
    for i in range(1, n):
        children[prev_room[i]].append(i)

    fact = [1] * (n + 1)
    for i in range(1, n + 1):
        fact[i] = fact[i - 1] * i % MOD

    inv_fact = [1] * (n + 1)
    inv_fact[n] = pow(fact[n], MOD - 2, MOD)
    for i in range(n - 1, -1, -1):
        inv_fact[i] = inv_fact[i + 1] * (i + 1) % MOD

    size = [0] * n
    ways = [1] * n

    order = []
    stack = [(0, False)]
    while stack:
        node, processed = stack.pop()
        if processed:
            order.append(node)
        else:
            stack.append((node, True))
            for child in children[node]:
                stack.append((child, False))

    for node in order:
        size[node] = 1
        total_children = 0
        for child in children[node]:
            size[node] += size[child]
            total_children += size[child]
            ways[node] = ways[node] * ways[child] % MOD
            ways[node] = ways[node] * inv_fact[size[child]] % MOD
        ways[node] = ways[node] * fact[total_children] % MOD

    return ways[0]

Complexity analysis

Time: You visit each node exactly once in the DFS. Precomputing factorials and inverse factorials is O(n). The modular exponentiation for the inverse factorial table is O(log p) which is constant relative to n. Total: O(n).

Space: The children adjacency list, factorial arrays, size array, and ways array each use O(n). The iterative DFS stack is O(n) in the worst case (a linear chain). Total: O(n).

The Building Blocks

This problem combines several reusable pieces that show up across many tree and combinatorics problems.

1. Building a rooted tree from a parent array

Converting a prevRoom array into an adjacency list is a pattern you will see in many tree problems:

children = [[] for _ in range(n)]
for i in range(1, n):
    children[prev_room[i]].append(i)

This is O(n) and gives you fast access to a node's children. The same pattern works any time you are given parent pointers instead of an explicit tree structure.

2. Iterative post-order DFS for subtree aggregation

Using an iterative stack with a "processed" flag to simulate post-order traversal:

stack = [(root, False)]
while stack:
    node, processed = stack.pop()
    if processed:
        # handle node after children
    else:
        stack.append((node, True))
        for child in children[node]:
            stack.append((child, False))

This avoids recursion depth limits (Python's default is 1000) and processes children before parents, exactly what you need when computing subtree sizes or aggregating results bottom-up.

3. Precomputed factorials with modular inverse

For any problem that needs combinations, permutations, or multinomial coefficients modulo a prime:

fact[i] = fact[i-1] * i % MOD
inv_fact[n] = pow(fact[n], MOD - 2, MOD)
inv_fact[i] = inv_fact[i+1] * (i+1) % MOD

Precomputing both fact and inv_fact arrays lets you evaluate any binomial or multinomial coefficient in O(1). This is a fundamental technique in competitive programming that transfers directly to interview problems involving counting under modular arithmetic.

Edge cases to watch for

• Single room: prevRoom = [-1]. Only one room, one way to build it. Return 1.
• Linear chain: prevRoom = [-1, 0, 1, 2, ...]. Every node has at most one child. Only one valid topological ordering exists. The multinomial coefficients are all 1.
• Star graph: prevRoom = [-1, 0, 0, 0, ...]. Every room depends only on room 0. The answer is (n-1)! because after building room 0, you can build the remaining n-1 rooms in any order.
• Large n: The problem allows up to n = 10^5. Recursive DFS will hit Python's stack limit. Use the iterative approach shown above.
• Modular arithmetic: Forgetting to take the modulus at each multiplication step will cause integer overflow in languages with fixed-width integers. In Python you will not overflow, but you will get the wrong answer if you forget the final modulus.

From understanding to recall

This problem ties together tree traversal, subtree size computation, and modular combinatorics. Each piece is manageable on its own, but combining them under interview pressure is where things get tricky. You need to remember the multinomial coefficient formula, know how to precompute factorials and modular inverses, and wire it all together in a clean DFS.

Spaced repetition bridges the gap between "I understand the approach" and "I can write it from memory." You drill the factorial precomputation pattern, the iterative post-order DFS, and the multinomial interleaving formula as separate building blocks. After a few rounds of practice at increasing intervals, you will not need to re-derive any of it. The code just flows.

Related posts

• Course Schedule - The foundation for topological ordering problems, teaching cycle detection and dependency graph traversal
• Binary Tree Maximum Path Sum - Another hard tree problem using post-order DFS where you aggregate child results upward
• Unique Binary Search Trees - Counting tree structures with combinatorics (Catalan numbers), showing how math and DP combine on tree counting problems

CodeBricks breaks Count Ways to Build Rooms into its tree-traversal, subtree-aggregation, and modular-combinatorics building blocks, then drills them independently with spaced repetition. You type each piece from scratch until the pattern is automatic. When a tree combinatorics problem shows up in your interview, you do not hesitate. You just write it.