LeetCode 1916 - Count Ways to Build Rooms in an Ant Colony

The problem asks us to count the number of valid ways to sequentially build all rooms in an ant colony, given a dependency tree that dictates which rooms must be built before others. Each room i has a predecessor prevRoom[i] that must be built before i.

leetcodehardarraymathdynamic-programmingtreedepth-first-searchgraph-theorytopological-sortcombinatorics

LeetCode Problem 1916

Difficulty: 🔴 Hard
Topics: Array, Math, Dynamic Programming, Tree, Depth-First Search, Graph Theory, Topological Sort, Combinatorics

Solution

Problem Understanding

The problem asks us to count the number of valid ways to sequentially build all rooms in an ant colony, given a dependency tree that dictates which rooms must be built before others. Each room i has a predecessor prevRoom[i] that must be built before i. Room 0 is always built first. The input is a list of integers, prevRoom, where each element either points to the parent room or is -1 for the root. The output is a single integer representing the number of valid sequences to build all rooms modulo 10^9 + 7.

Constraints tell us that n can be as large as 10^5, so a brute-force solution generating all permutations is infeasible. We also know the input forms a valid tree rooted at room 0, which guarantees there are no cycles and all rooms are reachable. Edge cases include chains of rooms (linear dependencies), trees with high branching factor, and minimal input where n = 2.

The key challenge is that multiple rooms can often be built in different orders if they share the same parent, which is a combinatorial counting problem over the tree structure.

Approaches

Brute Force

A brute-force approach would attempt to generate all possible sequences that respect the dependency constraints. This can be done by recursively exploring all rooms whose parent has already been built. While this approach is correct in principle, it is exponentially slow because the number of sequences grows factorially with the number of rooms, making it impossible to handle n = 10^5.

Optimal Approach

The key insight is that the problem can be reduced to counting the number of ways to interleave subtrees in a rooted tree. For each node, we can recursively compute the number of ways to build its subtree, then multiply by the multinomial coefficient representing the number of ways to arrange its child subtrees. This requires:

  1. Building an adjacency list representation of the tree.
  2. Using DFS to compute subtree sizes and the number of valid build sequences for each subtree.
  3. Applying combinatorics: If a node has k children with subtree sizes [s1, s2, ..., sk], the number of ways to interleave them is (s1 + s2 + ... + sk)! / (s1! * s2! * ... * sk!).

We precompute factorials modulo 10^9 + 7 to make calculations efficient.

Approach Time Complexity Space Complexity Notes
Brute Force O(n!) O(n) Recursively generate all sequences; infeasible for large n
Optimal O(n log n) O(n) DFS + combinatorics using precomputed factorials; handles large n efficiently

Algorithm Walkthrough

  1. Precompute Factorials and Inverses: We compute fact[i] = i! % MOD and inv[i] = modular inverse of fact[i] using Fermat's little theorem, to allow O(1) computation of binomial coefficients modulo 10^9 + 7.
  2. Build Tree: Convert prevRoom into an adjacency list tree where each node stores a list of its children. This will let us traverse the tree efficiently.
  3. DFS Function: Define a recursive DFS function that returns (ways, size) for a subtree rooted at node u:
  • size is the total number of nodes in the subtree.
  • ways is the number of valid build sequences for that subtree.
  1. Recursive Step:
  • Initialize total_ways = 1 and total_size = 0.
  • For each child v of u, recursively compute (child_ways, child_size).
  • Multiply total_ways by child_ways.
  • Multiply by the multinomial coefficient C(total_size + child_size, child_size) using precomputed factorials.
  • Increment total_size by child_size.
  1. Return Result: After DFS from the root 0, return ways as the final answer.

Why it works: Each node computes the number of valid sequences of building its subtree using the product of its children's ways and interleaving factor. Because the tree has no cycles, this recursion correctly counts all possible sequences without overcounting.

Python Solution

from typing import List

MOD = 10**9 + 7

class Solution:

    def waysToBuildRooms(self, prevRoom: List[int]) -> int:
        n = len(prevRoom)
        
        # Precompute factorials and modular inverses
        fact = [1] * (n + 1)
        inv = [1] * (n + 1)
        for i in range(2, n + 1):
            fact[i] = fact[i - 1] * i % MOD
        inv[n] = pow(fact[n], MOD - 2, MOD)
        for i in range(n - 1, 0, -1):
            inv[i] = inv[i + 1] * (i + 1) % MOD
        
        # Build tree
        tree = [[] for _ in range(n)]
        for child, parent in enumerate(prevRoom):
            if parent != -1:
                tree[parent].append(child)
        
        # DFS to count ways
        def dfs(node: int) -> (int, int):
            total_ways, total_size = 1, 0
            for child in tree[node]:
                child_ways, child_size = dfs(child)
                total_ways = total_ways * child_ways % MOD
                total_ways = total_ways * fact[total_size + child_size] % MOD
                total_ways = total_ways * inv[total_size] % MOD
                total_ways = total_ways * inv[child_size] % MOD
                total_size += child_size
            return total_ways, total_size + 1
        
        result, _ = dfs(0)
        return result

This implementation follows the steps exactly. Factorials and modular inverses allow fast computation of binomial coefficients. The DFS accumulates subtree sizes and ways, combining them using multinomial coefficients.

Go Solution

package main

func waysToBuildRooms(prevRoom []int) int {
    const MOD int = 1e9 + 7
    n := len(prevRoom)
    
    // Precompute factorials and inverses
    fact := make([]int, n+1)
    inv := make([]int, n+1)
    fact[0] = 1
    for i := 1; i <= n; i++ {
        fact[i] = fact[i-1] * i % MOD
    }
    inv[n] = modPow(fact[n], MOD-2, MOD)
    for i := n - 1; i >= 0; i-- {
        inv[i] = inv[i+1] * (i+1) % MOD
    }
    
    // Build tree
    tree := make([][]int, n)
    for i := 0; i < n; i++ {
        if prevRoom[i] != -1 {
            tree[prevRoom[i]] = append(tree[prevRoom[i]], i)
        }
    }
    
    var dfs func(int) (int, int)
    dfs = func(node int) (int, int) {
        totalWays, totalSize := 1, 0
        for _, child := range tree[node] {
            childWays, childSize := dfs(child)
            totalWays = totalWays * childWays % MOD
            totalWays = totalWays * fact[totalSize+childSize] % MOD
            totalWays = totalWays * inv[totalSize] % MOD
            totalWays = totalWays * inv[childSize] % MOD
            totalSize += childSize
        }
        return totalWays, totalSize + 1
    }
    
    result, _ := dfs(0)
    return result
}

func modPow(a, b, mod int) int {
    res := 1
    for b > 0 {
        if b%2 == 1 {
            res = res * a % mod
        }
        a = a * a % mod
        b /= 2
    }
    return res
}

The Go implementation is similar but uses modPow for modular inverse calculation. The rest follows the same DFS and combinatorial logic. Go slices handle dynamic child lists efficiently.

Worked Examples

Example 1: prevRoom = [-1,0,1]

Node Children DFS returns (ways, size)
2 [] (1, 1)
1 [2] (1 * 1 * 1! / (0! * 1!) = 1, size=2)
0 [1] (1 * 1 * 1! / (0! * 2!) = 1, size=3)

Result: 1

Example 2: prevRoom = [-1,0,0,1,2]

Node Children DFS returns
3 [] (1,1)
1 [3] (1 * 1 * 1! / (0!*1!) =1, size=2)