Facebook Pixel

1916. Count Ways to Build Rooms in an Ant Colony

Problem Description

You're building n rooms numbered from 0 to n-1 for an ant colony. You're given an array prevRoom where prevRoom[i] tells you which room must be built before room i. Room 0 is already built (so prevRoom[0] = -1), and once you build a room, it's directly connected to its previous room.

Key constraints:

  • You can only build one room at a time
  • To build room i, room prevRoom[i] must already be built
  • You can only travel between rooms that are already built and connected
  • The rooms form a tree structure with room 0 as the root

Task: Count how many different orders you can build all the rooms. Return the answer modulo 10^9 + 7.

Example visualization: If prevRoom = [-1, 0, 0, 1], the structure looks like:

    0 (already built)
   / \
  1   2
 /
3

Valid build orders include:

  • Build 1, then 3, then 2
  • Build 1, then 2, then 3
  • Build 2, then 1, then 3

The solution uses a recursive approach with combinatorics. For each subtree, it calculates how many ways to interleave the building orders of its child subtrees using the binomial coefficient C(n+m, m), where n and m are the sizes of subtrees being combined. This accounts for all valid orderings while respecting the parent-child dependencies in the tree.

Quick Interview Experience
Help others by sharing your interview experience
Have you seen this problem before?

Intuition

Let's think about this problem step by step. Since we have a tree structure with dependencies, we need to build parent rooms before their children. This naturally suggests a recursive approach starting from the root.

Consider a simple case: if room 0 has two children (rooms 1 and 2), we can build either room 1 first or room 2 first. Both are valid since their parent (room 0) is already built.

Now, what if room 1 has its own subtree with 3 rooms total, and room 2 has a subtree with 2 rooms total? We need to figure out how many ways we can interleave the building of these two subtrees while maintaining their internal ordering.

This is where combinatorics comes in. If we have two sequences that must maintain their internal order, the number of ways to interleave them is the binomial coefficient C(n+m, m) where n and m are the lengths of the sequences. This is because we're choosing m positions out of n+m total positions for one sequence, and the other sequence fills the remaining positions.

For example, if subtree 1 has rooms [A, B] and subtree 2 has rooms [X, Y, Z], we can interleave them as:

  • A, B, X, Y, Z
  • A, X, B, Y, Z
  • A, X, Y, B, Z
  • X, A, B, Y, Z
  • ... and so on

The total number of such interleavings is C(2+3, 3) = C(5, 3) = 10.

When a node has multiple children, we process them one by one. First, we calculate the ways to build the first child's subtree. Then for each subsequent child, we multiply our current result by the number of ways to interleave that child's subtree with all previously processed subtrees. This multiplication principle works because the choices for each interleaving are independent.

The recursive function returns the size of each subtree (including the root), which we need to calculate the correct binomial coefficients when combining subtrees at higher levels of the tree.

Learn more about Tree, Graph, Topological Sort, Math, Dynamic Programming and Combinatorics patterns.

Solution Approach

The implementation uses a depth-first search (DFS) approach with combinatorial mathematics to count the valid building orders.

Data Structure Setup:

  • ingoing: A dictionary mapping each room to its parent (though not used in the main logic)
  • outgoing: A dictionary mapping each room to a set of its children
  • These dictionaries represent the tree structure as an adjacency list

Building the Tree Representation:

for i in range(1, len(prevRoom)):
    ingoing[i].add(prevRoom[i])
    outgoing[prevRoom[i]].add(i)

This loop constructs the parent-child relationships from the prevRoom array.

Recursive DFS Function: The recurse(i) function processes each node and returns the total number of nodes in its subtree (including itself).

For leaf nodes (rooms with no children):

if len(outgoing[i]) == 0:
    return 1

For internal nodes, we iterate through all children:

nodes_in_tree = 0
for v in outgoing[i]:
    cn = recurse(v)  # Get size of child's subtree
    if nodes_in_tree != 0:
        ans[0] *= comb(nodes_in_tree + cn, cn)
        ans[0] %= modulo
    nodes_in_tree += cn

Key Algorithm - Combining Subtrees: When we have processed some children (with nodes_in_tree total nodes) and encounter a new child subtree (with cn nodes), we multiply our current answer by C(nodes_in_tree + cn, cn).

This binomial coefficient represents the number of ways to interleave:

  • The nodes_in_tree nodes from previously processed subtrees
  • The cn nodes from the current child's subtree

The multiplication works because:

  1. Each subtree internally has its own number of valid orderings (calculated recursively)
  2. The ways to interleave different subtrees are independent choices
  3. By the multiplication principle, total ways = product of all these choices

Modular Arithmetic: Since the answer can be very large, we take modulo 10^9 + 7 after each multiplication to prevent overflow.

Final Result: The function starts with ans = [1] (using a list to maintain reference in nested function) and returns ans[0] % modulo as the final count of valid building orders.

Ready to land your dream job?

Unlock your dream job with a 3-minute evaluator for a personalized learning plan!

Start Evaluator

Example Walkthrough

Let's walk through the example where prevRoom = [-1, 0, 0, 1].

First, let's visualize the tree structure:

    0 (already built)
   / \
  1   2
 /
3

Step 1: Build Tree Representation

  • outgoing = {0: {1, 2}, 1: {3}, 2: {}, 3: {}}
  • Room 0 has children 1 and 2
  • Room 1 has child 3
  • Rooms 2 and 3 are leaves

Step 2: Start DFS from Room 0

  • Call recurse(0)
  • Room 0 has 2 children: 1 and 2
  • Initialize nodes_in_tree = 0, ans = [1]

Step 3: Process First Child (Room 1)

  • Call recurse(1)

    • Room 1 has 1 child: 3
    • Call recurse(3)
      • Room 3 is a leaf, returns 1
    • For room 1: nodes_in_tree = 0 initially
    • After processing child 3: nodes_in_tree = 1
    • Room 1's subtree has 2 nodes total (itself + room 3)
    • Returns 2
  • Back in recurse(0):

    • Child 1's subtree has 2 nodes
    • Since nodes_in_tree = 0 (first child), no multiplication needed
    • Update nodes_in_tree = 2

Step 4: Process Second Child (Room 2)

  • Call recurse(2)

    • Room 2 is a leaf, returns 1
  • Back in recurse(0):

    • Child 2's subtree has 1 node
    • Now we need to interleave:
      • 2 nodes from room 1's subtree (rooms 1 and 3)
      • 1 node from room 2's subtree (room 2)
    • Calculate: C(2+1, 1) = C(3, 1) = 3
    • Update answer: ans = 1 × 3 = 3

Step 5: Final Result The function returns 3, meaning there are 3 valid building orders:

  1. Build 1 → Build 3 → Build 2

    • After building 1, we can access and build 3
    • Finally build 2
  2. Build 1 → Build 2 → Build 3

    • Build 1 first, then 2
    • Finally build 3 (accessible through 1)
  3. Build 2 → Build 1 → Build 3

    • Build 2 first
    • Then build 1, finally 3

The key insight: The binomial coefficient C(3, 1) = 3 tells us there are 3 ways to choose 1 position out of 3 total positions for room 2, while rooms 1 and 3 must maintain their order (1 before 3).

Solution Implementation

1from typing import List
2from collections import defaultdict
3from math import comb
4
5class Solution:
6    def waysToBuildRooms(self, prevRoom: List[int]) -> int:
7        """
8        Calculate the number of ways to build rooms given their dependencies.
9      
10        Args:
11            prevRoom: List where prevRoom[i] is the room that must be built before room i
12      
13        Returns:
14            Number of valid build orders modulo 10^9 + 7
15        """
16        MOD = 10**9 + 7
17      
18        # Build adjacency lists for the tree structure
19        # ingoing[i] contains the parent of room i (not actually used in computation)
20        # outgoing[i] contains all children of room i
21        ingoing = defaultdict(set)
22        outgoing = defaultdict(set)
23      
24        # Construct the tree edges (room 0 is the root with no parent)
25        for room_id in range(1, len(prevRoom)):
26            ingoing[room_id].add(prevRoom[room_id])
27            outgoing[prevRoom[room_id]].add(room_id)
28      
29        # Use list to store answer to allow modification in nested function
30        result = [1]
31      
32        def count_subtree_arrangements(room: int) -> int:
33            """
34            Recursively count arrangements for the subtree rooted at 'room'.
35          
36            Args:
37                room: Current room being processed
38          
39            Returns:
40                Total number of nodes in the subtree rooted at 'room'
41            """
42            # Base case: leaf node (room with no children)
43            if len(outgoing[room]) == 0:
44                return 1
45          
46            total_nodes_in_subtrees = 0
47          
48            # Process each child subtree
49            for child_room in outgoing[room]:
50                # Get the size of child's subtree
51                child_subtree_size = count_subtree_arrangements(child_room)
52              
53                # Calculate combinations for merging this subtree with previous ones
54                # We need to choose positions for the new subtree nodes among all positions
55                if total_nodes_in_subtrees != 0:
56                    # Number of ways to interleave the sequences
57                    result[0] *= comb(total_nodes_in_subtrees + child_subtree_size, child_subtree_size)
58                    result[0] %= MOD
59              
60                total_nodes_in_subtrees += child_subtree_size
61          
62            # Return total nodes in this subtree (including current room)
63            return total_nodes_in_subtrees + 1
64      
65        # Start DFS from root (room 0)
66        count_subtree_arrangements(0)
67      
68        return result[0] % MOD
69
1import java.util.*;
2
3class Solution {
4    private static final int MOD = 1000000007;
5    private long result;
6    private Map<Integer, Set<Integer>> outgoing;
7  
8    public int waysToBuildRooms(int[] prevRoom) {
9        /**
10         * Calculate the number of ways to build rooms given their dependencies.
11         * 
12         * @param prevRoom Array where prevRoom[i] is the room that must be built before room i
13         * @return Number of valid build orders modulo 10^9 + 7
14         */
15      
16        // Build adjacency lists for the tree structure
17        // outgoing[i] contains all children of room i
18        outgoing = new HashMap<>();
19      
20        // Initialize the outgoing map for all rooms
21        for (int i = 0; i < prevRoom.length; i++) {
22            outgoing.put(i, new HashSet<>());
23        }
24      
25        // Construct the tree edges (room 0 is the root with no parent)
26        for (int roomId = 1; roomId < prevRoom.length; roomId++) {
27            outgoing.get(prevRoom[roomId]).add(roomId);
28        }
29      
30        // Initialize result
31        result = 1;
32      
33        // Precompute factorials and inverse factorials for combination calculation
34        long[] factorial = new long[prevRoom.length + 1];
35        long[] invFactorial = new long[prevRoom.length + 1];
36        factorial[0] = 1;
37        for (int i = 1; i <= prevRoom.length; i++) {
38            factorial[i] = (factorial[i - 1] * i) % MOD;
39        }
40        invFactorial[prevRoom.length] = modInverse(factorial[prevRoom.length], MOD);
41        for (int i = prevRoom.length - 1; i >= 0; i--) {
42            invFactorial[i] = (invFactorial[i + 1] * (i + 1)) % MOD;
43        }
44      
45        // Start DFS from root (room 0)
46        countSubtreeArrangements(0, factorial, invFactorial);
47      
48        return (int)(result % MOD);
49    }
50  
51    private int countSubtreeArrangements(int room, long[] factorial, long[] invFactorial) {
52        /**
53         * Recursively count arrangements for the subtree rooted at 'room'.
54         * 
55         * @param room Current room being processed
56         * @param factorial Precomputed factorial values
57         * @param invFactorial Precomputed inverse factorial values
58         * @return Total number of nodes in the subtree rooted at 'room'
59         */
60      
61        // Base case: leaf node (room with no children)
62        if (outgoing.get(room).isEmpty()) {
63            return 1;
64        }
65      
66        int totalNodesInSubtrees = 0;
67      
68        // Process each child subtree
69        for (int childRoom : outgoing.get(room)) {
70            // Get the size of child's subtree
71            int childSubtreeSize = countSubtreeArrangements(childRoom, factorial, invFactorial);
72          
73            // Calculate combinations for merging this subtree with previous ones
74            // We need to choose positions for the new subtree nodes among all positions
75            if (totalNodesInSubtrees != 0) {
76                // Number of ways to interleave the sequences
77                // C(totalNodesInSubtrees + childSubtreeSize, childSubtreeSize)
78                long combination = factorial[totalNodesInSubtrees + childSubtreeSize];
79                combination = (combination * invFactorial[childSubtreeSize]) % MOD;
80                combination = (combination * invFactorial[totalNodesInSubtrees]) % MOD;
81                result = (result * combination) % MOD;
82            }
83          
84            totalNodesInSubtrees += childSubtreeSize;
85        }
86      
87        // Return total nodes in this subtree (including current room)
88        return totalNodesInSubtrees + 1;
89    }
90  
91    private long modInverse(long a, long m) {
92        /**
93         * Calculate modular multiplicative inverse using Fermat's little theorem.
94         * Since m is prime, a^(-1) ≡ a^(m-2) (mod m)
95         * 
96         * @param a The number to find inverse of
97         * @param m The modulus
98         * @return Modular inverse of a
99         */
100        return modPow(a, m - 2, m);
101    }
102  
103    private long modPow(long base, long exp, long mod) {
104        /**
105         * Calculate (base^exp) % mod using binary exponentiation.
106         * 
107         * @param base Base number
108         * @param exp Exponent
109         * @param mod Modulus
110         * @return (base^exp) % mod
111         */
112        long result = 1;
113        base %= mod;
114        while (exp > 0) {
115            if ((exp & 1) == 1) {
116                result = (result * base) % mod;
117            }
118            base = (base * base) % mod;
119            exp >>= 1;
120        }
121        return result;
122    }
123}
124
1class Solution {
2private:
3    const int MOD = 1000000007;
4  
5    // Precompute factorials and inverse factorials for combinations
6    vector<long long> factorial;
7    vector<long long> invFactorial;
8  
9    // Fast modular exponentiation for computing modular inverse
10    long long modPow(long long base, long long exp, long long mod) {
11        long long result = 1;
12        while (exp > 0) {
13            if (exp % 2 == 1) {
14                result = (result * base) % mod;
15            }
16            base = (base * base) % mod;
17            exp /= 2;
18        }
19        return result;
20    }
21  
22    // Calculate nCr modulo MOD using precomputed factorials
23    long long combination(int n, int r) {
24        if (r > n || r < 0) return 0;
25        return (factorial[n] * invFactorial[r] % MOD) * invFactorial[n - r] % MOD;
26    }
27  
28    // Precompute factorials and their modular inverses
29    void precomputeFactorials(int n) {
30        factorial.resize(n + 1);
31        invFactorial.resize(n + 1);
32      
33        factorial[0] = 1;
34        for (int i = 1; i <= n; i++) {
35            factorial[i] = (factorial[i - 1] * i) % MOD;
36        }
37      
38        invFactorial[n] = modPow(factorial[n], MOD - 2, MOD);
39        for (int i = n - 1; i >= 0; i--) {
40            invFactorial[i] = (invFactorial[i + 1] * (i + 1)) % MOD;
41        }
42    }
43  
44public:
45    int waysToBuildRooms(vector<int>& prevRoom) {
46        int n = prevRoom.size();
47      
48        // Precompute factorials for combination calculations
49        precomputeFactorials(n);
50      
51        // Build adjacency list for the tree structure
52        // outgoing[i] contains all children of room i
53        vector<vector<int>> outgoing(n);
54      
55        // Construct the tree edges (room 0 is the root with no parent)
56        for (int roomId = 1; roomId < n; roomId++) {
57            outgoing[prevRoom[roomId]].push_back(roomId);
58        }
59      
60        // Store the result as we traverse
61        long long result = 1;
62      
63        // Lambda function to recursively count arrangements for subtree rooted at 'room'
64        function<int(int)> countSubtreeArrangements = [&](int room) -> int {
65            // Base case: leaf node (room with no children)
66            if (outgoing[room].empty()) {
67                return 1;
68            }
69          
70            int totalNodesInSubtrees = 0;
71          
72            // Process each child subtree
73            for (int childRoom : outgoing[room]) {
74                // Get the size of child's subtree
75                int childSubtreeSize = countSubtreeArrangements(childRoom);
76              
77                // Calculate combinations for merging this subtree with previous ones
78                // We need to choose positions for the new subtree nodes among all positions
79                if (totalNodesInSubtrees != 0) {
80                    // Number of ways to interleave the sequences
81                    result = (result * combination(totalNodesInSubtrees + childSubtreeSize, 
82                                                   childSubtreeSize)) % MOD;
83                }
84              
85                totalNodesInSubtrees += childSubtreeSize;
86            }
87          
88            // Return total nodes in this subtree (including current room)
89            return totalNodesInSubtrees + 1;
90        };
91      
92        // Start DFS from root (room 0)
93        countSubtreeArrangements(0);
94      
95        return result;
96    }
97};
98
1function waysToBuildRooms(prevRoom: number[]): number {
2    /**
3     * Calculate the number of ways to build rooms given their dependencies.
4     * 
5     * @param prevRoom - Array where prevRoom[i] is the room that must be built before room i
6     * @returns Number of valid build orders modulo 10^9 + 7
7     */
8    const MOD = 10 ** 9 + 7;
9  
10    // Build adjacency lists for the tree structure
11    // outgoing[i] contains all children of room i
12    const outgoing: Map<number, Set<number>> = new Map();
13  
14    // Initialize the map with empty sets for all rooms
15    for (let i = 0; i < prevRoom.length; i++) {
16        outgoing.set(i, new Set());
17    }
18  
19    // Construct the tree edges (room 0 is the root with no parent)
20    for (let roomId = 1; roomId < prevRoom.length; roomId++) {
21        outgoing.get(prevRoom[roomId])!.add(roomId);
22    }
23  
24    // Store the result as a mutable reference
25    let result = 1;
26  
27    /**
28     * Calculate binomial coefficient (n choose k)
29     * @param n - Total number of items
30     * @param k - Number of items to choose
31     * @returns The binomial coefficient modulo MOD
32     */
33    function binomialCoefficient(n: number, k: number): number {
34        if (k > n || k < 0) return 0;
35        if (k === 0 || k === n) return 1;
36      
37        // Use Pascal's triangle approach with modular arithmetic
38        let numerator = 1;
39        let denominator = 1;
40      
41        for (let i = 0; i < k; i++) {
42            numerator = (numerator * (n - i)) % MOD;
43            denominator = (denominator * (i + 1)) % MOD;
44        }
45      
46        // Calculate modular inverse of denominator
47        return (numerator * modularInverse(denominator, MOD)) % MOD;
48    }
49  
50    /**
51     * Calculate modular inverse using Fermat's little theorem
52     * @param a - Number to find inverse of
53     * @param mod - Modulus value
54     * @returns Modular inverse of a
55     */
56    function modularInverse(a: number, mod: number): number {
57        return modularPower(a, mod - 2, mod);
58    }
59  
60    /**
61     * Calculate (base^exp) % mod efficiently
62     * @param base - Base number
63     * @param exp - Exponent
64     * @param mod - Modulus value
65     * @returns Result of modular exponentiation
66     */
67    function modularPower(base: number, exp: number, mod: number): number {
68        let result = 1;
69        base %= mod;
70      
71        while (exp > 0) {
72            if (exp & 1) {
73                result = (result * base) % mod;
74            }
75            base = (base * base) % mod;
76            exp >>= 1;
77        }
78      
79        return result;
80    }
81  
82    /**
83     * Recursively count arrangements for the subtree rooted at 'room'.
84     * 
85     * @param room - Current room being processed
86     * @returns Total number of nodes in the subtree rooted at 'room'
87     */
88    function countSubtreeArrangements(room: number): number {
89        const children = outgoing.get(room)!;
90      
91        // Base case: leaf node (room with no children)
92        if (children.size === 0) {
93            return 1;
94        }
95      
96        let totalNodesInSubtrees = 0;
97      
98        // Process each child subtree
99        for (const childRoom of children) {
100            // Get the size of child's subtree
101            const childSubtreeSize = countSubtreeArrangements(childRoom);
102          
103            // Calculate combinations for merging this subtree with previous ones
104            // We need to choose positions for the new subtree nodes among all positions
105            if (totalNodesInSubtrees !== 0) {
106                // Number of ways to interleave the sequences
107                const combinations = binomialCoefficient(
108                    totalNodesInSubtrees + childSubtreeSize, 
109                    childSubtreeSize
110                );
111                result = (result * combinations) % MOD;
112            }
113          
114            totalNodesInSubtrees += childSubtreeSize;
115        }
116      
117        // Return total nodes in this subtree (including current room)
118        return totalNodesInSubtrees + 1;
119    }
120  
121    // Start DFS from root (room 0)
122    countSubtreeArrangements(0);
123  
124    return result % MOD;
125}
126

Time and Space Complexity

Time Complexity: O(n²)

The algorithm performs a depth-first traversal of the tree starting from room 0. For each node visited:

  • The recursive function recurse() is called exactly once for each of the n rooms
  • At each node with children, the algorithm computes combinations using comb(nodes_in_tree + cn, cn)
  • The comb() function has a time complexity of O(n) in the worst case when computing binomial coefficients
  • In the worst case (e.g., a linear tree structure), the sum of all combination computations across all nodes can reach O(n²)

Therefore, the overall time complexity is O(n²).

Space Complexity: O(n)

The space usage consists of:

  • ingoing dictionary: stores at most n-1 entries (one parent per room except room 0), using O(n) space
  • outgoing dictionary: stores all parent-child relationships, totaling n-1 edges, using O(n) space
  • Recursive call stack: in the worst case (linear tree), the recursion depth can be O(n)
  • ans list: stores a single value, using O(1) space

The dominant factor is the recursive call stack and the dictionaries, resulting in O(n) space complexity.

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Integer Overflow in Combinatorial Calculations

The most critical pitfall in this solution is handling large combinatorial values. The comb(n, k) function can produce extremely large numbers that exceed standard integer limits, especially when dealing with large tree structures.

The Problem:

# This can cause overflow or incorrect results
result[0] *= comb(total_nodes_in_subtrees + child_subtree_size, child_subtree_size)
result[0] %= MOD

When comb() returns a very large number, multiplying it with result[0] before taking modulo can cause overflow issues in some implementations or languages.

Solution: Implement a modular arithmetic version of the combination function:

def comb_mod(n, k, mod):
    """Calculate C(n,k) % mod using modular arithmetic"""
    if k > n or k < 0:
        return 0
    if k == 0 or k == n:
        return 1
  
    # Calculate using modular arithmetic to prevent overflow
    numerator = 1
    denominator = 1
  
    for i in range(min(k, n - k)):
        numerator = (numerator * (n - i)) % mod
        denominator = (denominator * (i + 1)) % mod
  
    # Use modular inverse for division
    return (numerator * pow(denominator, mod - 2, mod)) % mod

2. Incorrect Tree Construction

Another common mistake is misunderstanding the parent-child relationships when building the adjacency list.

The Problem:

# Wrong direction - treating prevRoom[i] as child of i
outgoing[i].add(prevRoom[i])  # INCORRECT!

Solution: Always remember that prevRoom[i] is the parent of room i, not the other way around:

# Correct: prevRoom[i] is the parent, so i is a child of prevRoom[i]
outgoing[prevRoom[i]].add(i)

3. Not Handling Single Child Optimization

The current code has an unnecessary check that could be simplified:

The Problem:

if total_nodes_in_subtrees != 0:
    result[0] *= comb(total_nodes_in_subtrees + child_subtree_size, child_subtree_size)

This check prevents multiplication on the first child, but comb(0 + cn, cn) equals 1 anyway, making the check redundant.

Solution: Remove the conditional for cleaner code:

# Always multiply - when total_nodes_in_subtrees is 0, comb(cn, cn) = 1
result[0] = (result[0] * comb(total_nodes_in_subtrees + child_subtree_size, 
                              child_subtree_size)) % MOD

4. Using Mutable Default Arguments

While not present in this code, a common Python pitfall when using defaultdict is accidentally sharing mutable defaults across function calls.

Solution: Always create new instances inside the function rather than using default parameters:

# Good practice
def waysToBuildRooms(self, prevRoom: List[int]) -> int:
    ingoing = defaultdict(set)  # Fresh instance each call
    outgoing = defaultdict(set)
Discover Your Strengths and Weaknesses: Take Our 3-Minute Quiz to Tailor Your Study Plan:

The three-steps of Depth First Search are:

  1. Identify states;
  2. Draw the state-space tree;
  3. DFS on the state-space tree.

Recommended Readings

Want a Structured Path to Master System Design Too? Don’t Miss This!

Load More