Count Ways to Build Rooms in an Ant Colony solution and explanation
3 more videos on this problem
LeetCode 1916. Count Ways to Build Rooms in an Ant ColonyHappy Coding · 34:34 · 3K views
1916. Count Ways to Build Rooms in an Ant Colony (Leetcode Hard)Programming Live with Larry · 39:08 · 1K views
【每日一题】1916. Count Ways to Build Rooms in an Ant Colony, 6/30/2021Huifeng Guan · 26:13 · 748 views
LeetCode 1916: Count Ways to Build Rooms in an Ant ColonyAlitaCode · 31:34 · 42 views
The problem
You are an ant tasked with adding n new rooms numbered 0 to n-1 to your colony. You are given the expansion plan as a 0-indexed integer array of length n, prevRoom, where prevRoom[i] indicates that you must build room prevRoom[i] before building room i, and these two rooms must be connected directly. Room 0 is already built, so prevRoom[0] = -1. The expansion plan is given such that once all the rooms are built, every room will be reachable from room 0.
You can only build one room at a time, and you can travel freely between rooms you have already built only if they are connected. You can choose to build any room as long as its previous room is already built.
Return the number of different orders you can build all the rooms in. Since the answer may be large, return it modulo 109 + 7.
Example 1:
Input: prevRoom = [-1,0,1] Output: 1 Explanation: There is only one way to build the additional rooms: 0 → 1 → 2
Example 2:
Input: prevRoom = [-1,0,0,1,2] Output: 6 Explanation: The 6 ways are: 0 → 1 → 3 → 2 → 4 0 → 2 → 4 → 1 → 3 0 → 1 → 2 → 3 → 4 0 → 1 → 2 → 4 → 3 0 → 2 → 1 → 3 → 4 0 → 2 → 1 → 4 → 3
Constraints:
n == prevRoom.length2 <= n <= 105prevRoom[0] == -10 <= prevRoom[i] < nfor all1 <= i < n- Every room is reachable from room
0once all the rooms are built.
Approach overview
Problem Overview: You are given a tree where each node represents a room and prevRoom[i] indicates which room must be built before room i. Room 0 is already built. The task is to count how many different valid orders exist to build all rooms while respecting these parent dependencies.
Approach 1: DFS and Combinatorics (O(n) time, O(n) space)
Treat the dependency structure as a rooted tree. Each subtree can be built independently once its parent is finished. During a DFS, compute the size of each subtree and the number of valid ways to arrange nodes inside it. When combining children, use combinatorics: if subtrees have sizes s1, s2, ..., the number of ways to interleave them is a multinomial coefficient. Precompute factorials and modular inverses to evaluate combinations under modulo. This approach leverages tree traversal and combinatorics to merge subtree results efficiently.
Approach 2: Dynamic Programming with Tree Representation (O(n) time, O(n) space)
Build an adjacency list from the parent array and run a DFS that returns two values for each node: the number of nodes in its subtree and the number of valid build orders for that subtree. While processing children, accumulate subtree sizes and multiply the child arrangements with combination counts representing interleavings. The DP state essentially captures "ways to build this subtree" while respecting dependencies. This formulation fits naturally with dynamic programming on trees.
Approach 3: Combination Mathematics in Tree (O(n) time, O(n) space)
The mathematical perspective models the problem as counting valid topological orders of a rooted tree. For each node with children subtrees of sizes s1, s2, ..., you place their nodes among s1 + s2 + ... slots. The number of placements equals (s1+s2+...)! / (s1! * s2! * ...). Multiply this with the number of arrangements inside each subtree. Precompute factorials and inverse factorials once, then evaluate each combination in constant time.
Approach 4: Topological Sort with Dynamic Programming (O(n) time, O(n) space)
Interpret the dependency tree as a DAG and process nodes using a topological sort. Maintain DP values representing ways to construct partial subgraphs and combine them when prerequisites are satisfied. Although the structure is still a tree, framing it as a topological ordering problem helps when reasoning about dependency scheduling and incremental state transitions.
Recommended for interviews: DFS with combinatorics is the expected solution. It shows you understand how subtree independence leads to multinomial combinations. A brute-force enumeration of topological orders quickly becomes exponential, while the combinatorial tree DP reduces the problem to linear time with precomputed factorials.
Approach 1: Dynamic Programming with Tree Representation
Represent the problem as a tree with each node referring to a room, and edges between nodes pointing to prerequisite dependencies. Use a dynamic programming approach to count the number of permutations facilitated by the tree structure.
Consider breaking down the problem into subproblems based on subtrees and calculate the number of valid orders for each subtree.
from collections import defaultdict
MOD = 10**9 + 7
def dfs(node, graph, node_count, combo_cache):
total_ways = 1
child_size_sum = 0
for child in graph[node]:
child_ways, child_size = dfs(child, graph, node_count, combo_cache)
total_ways *= child_ways
total_ways %= MOD
total_ways *= combo_cache[child_size_sum + child_size][child_size]
total_ways %= MOD
child_size_sum += child_size
node_count[node] = 1 + child_size_sum
return total_ways, node_count[node]
def countPermutationsOfTree(prevRoom):
n = len(prevRoom)
graph = defaultdict(list)
for room in range(1, n):
graph[prevRoom[room]].append(room)
# Precompute combinations
combo_cache = [[1] * (n + 1) for _ in range(n + 1)]
for i in range(1, n + 1):
for j in range(1, i):
combo_cache[i][j] = (combo_cache[i - 1][j - 1] + combo_cache[i - 1][j]) % MOD
# Use DFS to calculate result
node_count = [0] * n
result, _ = dfs(0, graph, node_count, combo_cache)
return result
# Test
prevRoom = [-1, 0, 0, 1, 2]
print(countPermutationsOfTree(prevRoom))This code represents the rooms in a tree structure using a directed graph with nodes and edges indicating the prerequisite room. We use a depth-first search (DFS) approach to traverse the tree, computing permutations dynamically at each node (room) level using combination mathematics. A 2D list is precomputed for combination calculations using Pascal's Triangle logic.
Time Complexity: O(n^2), where n is the number of rooms, due to combinations precomputation. Space Complexity: O(n^2) for the combination table to store computed values.
Approach 2: Combination Mathematics in Tree
Use tree structure to identify dependent builds. Solve using DFS with combination calculations derived from subtree sizes. This approach merges combinatorial logic into tree traversal to efficiently calculate possible room-building sequences.
const MOD = 1000000007;
function combination(n, k, comboCache) {
if (k === 0 || n === k) return 1;
if (comboCache[n][k] !== null) return comboCache[n][k];
comboCache[n][k] = (combination(n - 1, k - 1, comboCache) + combination(n - 1, k, comboCache)) % MOD;
return comboCache[n][k];
}
function dfs(node, tree, nodeCount, comboCache) {
let totalWays = 1;
let sumChildrenSizes = 0;
for (const child of tree[node]) {
const [childWays, childSize] = dfs(child, tree, nodeCount, comboCache);
totalWays = (totalWays * childWays) % MOD;
totalWays = (totalWays * combination(sumChildrenSizes + childSize, childSize, comboCache)) % MOD;
sumChildrenSizes += childSize;
}
nodeCount[node] = 1 + sumChildrenSizes;
return [totalWays, nodeCount[node]];
}
function countWaysToBuildRooms(prevRoom) {
const n = prevRoom.length;
const tree = Array.from({ length: n }, () => []);
for (let i = 1; i < n; i++) {
tree[prevRoom[i]].push(i);
}
// Combination cache
const comboCache = Array.from({ length: n + 1 }, () => Array(n + 1).fill(null));
const nodeCount = Array(n).fill(0);
return dfs(0, tree, nodeCount, comboCache)[0];
}
// Test
console.log(countWaysToBuildRooms([-1, 0, 0, 1, 2]));This JavaScript function leverages a tree structure and employs a recursive DFS approach to compute room-building permutations. Combinatorial logic is embedded into the recursion to account for subtree permutation combinations. A cache optimizes combination computations based on Pascal’s Triangle principles.
Time Complexity: O(n^2), due to combination calculations. Space Complexity: O(n^2), used for the combination cache.
Approach 3: Approach 1: DFS and Combinatorics
In this approach, think of the colony structure as a tree rooted at room 0. For each node (room), we can choose its children in any order as long as we ensure the order dictated by prevRoom. We use Depth-First Search (DFS) to explore this tree, calculating the number of ways we can construct the subtree rooted at each node while considering its children placement order using combinatorial principles. We employ a comb(n, k) approach, where n is the total number of nodes and k are the children nodes, to determine different permutations of building sequences taking into account already built sequences of other branches.
from functools import lru_cache
MOD = 10**9 + 7
def fac(n):
result = 1
for i in range(2, n+1):
result = result * i % MOD
return result
def comb(n, k):
if k > n:
return 0
return fac(n) * pow(fac(k), MOD-2, MOD) % MOD * pow(fac(n-k), MOD-2, MOD) % MOD
class Solution:
def waysToBuildRooms(self, prevRoom):
n = len(prevRoom)
tree = [[] for _ in range(n)]
for i in range(1, n):
tree[prevRoom[i]].append(i)
def dfs(node):
way = 1
size = 0
for ch in tree[node]:
sub_way, sub_size = dfs(ch)
way *= comb(size + sub_size, sub_size) * sub_way % MOD
size += sub_size
way %= MOD
return way, size + 1
return dfs(0)[0]import java.util.ArrayList;
import java.util.List;
class Solution {
private static final int MOD = (int)1e9 + 7;
private long factorial(int n) {
long result = 1;
for (int i = 2; i <= n; i++) {
result = result * i % MOD;
}
return result;
}
private long comb(int n, int k) {
if (k > n) return 0;
return factorial(n) * modInverse(factorial(k), MOD) % MOD * modInverse(factorial(n-k), MOD) % MOD;
}
private long modInverse(long a, long m) {
return power(a, m - 2, m);
}
private long power(long x, long y, long p) {
long res = 1;
x = x % p;
while (y > 0) {
if (y % 2 == 1)
res = (res * x) % p;
y = y >> 1;
x = (x * x) % p;
}
return res;
}
public int waysToBuildRooms(int[] prevRoom) {
int n = prevRoom.length;
List<List<Integer>> tree = new ArrayList<>();
for (int i = 0; i < n; i++) tree.add(new ArrayList<>());
for (int i = 1; i < n; i++) tree.get(prevRoom[i]).add(i);
return (int) dfs(tree, 0)[0];
}
private long[] dfs(List<List<Integer>> tree, int node) {
long way = 1;
int size = 0;
for (int ch : tree.get(node)) {
long[] sub = dfs(tree, ch);
way = way * comb(size + (int) sub[1], (int) sub[1]) % MOD * sub[0] % MOD;
size += (int) sub[1];
}
return new long[]{way, size + 1};
}
}#include <vector>
using namespace std;
static const int MOD = 1e9 + 7;
long long power(long long x, long long y, long long p) {
long long res = 1;
x = x % p;
while (y > 0) {
if (y & 1)
res = (res * x) % p;
y = y >> 1;
x = (x * x) % p;
}
return res;
}
long long modInverse(long long a, long long m) {
return power(a, m - 2, m);
}
long long factorial(int n) {
long long result = 1;
for (int i = 2; i <= n; i++) {
result = result * i % MOD;
}
return result;
}
long long comb(int n, int k) {
if (k > n) return 0;
return factorial(n) * modInverse(factorial(k), MOD) % MOD * modInverse(factorial(n - k), MOD) % MOD;
}
class Solution {
public:
int waysToBuildRooms(vector<int>& prevRoom) {
int n = prevRoom.size();
vector<vector<int>> tree(n);
for (int i = 1; i < n; i++)
tree[prevRoom[i]].push_back(i);
return dfs(tree, 0).first;
}
private:
pair<long long int, int> dfs(vector<vector<int>>& tree, int node) {
long long way = 1;
int size = 0;
for (int child : tree[node]) {
auto sub = dfs(tree, child);
way = way * comb(size + sub.second, sub.second) % MOD * sub.first % MOD;
size += sub.second;
}
return {way, size + 1};
}
};const MOD = 1e9 + 7;
function factorial(n) {
let result = 1;
for (let i = 2; i <= n; i++) {
result = (result * i) % MOD;
}
return result;
}
function power(x, y, mod) {
let res = 1;
x = x % mod;
while (y > 0) {
if (y % 2 === 1) {
res = (res * x) % mod;
}
y = Math.floor(y / 2);
x = (x * x) % mod;
}
return res;
}
function modInverse(a, m) {
return power(a, m - 2, m);
}
function comb(n, k) {
if (k > n) return 0;
return (factorial(n) * modInverse(factorial(k), MOD) % MOD * modInverse(factorial(n - k), MOD) % MOD) % MOD;
}
function waysToBuildRooms(prevRoom) {
const n = prevRoom.length;
const tree = Array.from({ length: n }, () => []);
for (let i = 1; i < n; i++) {
tree[prevRoom[i]].push(i);
}
return dfs(0)[0];
function dfs(node) {
let way = 1;
let size = 0;
for (const child of tree[node]) {
const [subWay, subSize] = dfs(child);
way = way * comb(size + subSize, subSize) % MOD * subWay % MOD;
size += subSize;
}
return [way, size + 1];
}
}using System;
using System.Collections.Generic;
class Solution {
private static int MOD = (int)(1e9 + 7);
private long Factorial(int n) {
long result = 1;
for (int i = 2; i <= n; i++) {
result = result * i % MOD;
}
return result;
}
private long Comb(int n, int k) {
if (k > n) return 0;
return Factorial(n) * ModInverse(Factorial(k), MOD) % MOD * ModInverse(Factorial(n - k), MOD) % MOD;
}
private long ModInverse(long a, long m) {
return Power(a, m - 2, m);
}
private long Power(long x, long y, long p) {
long res = 1;
x = x % p;
while (y > 0) {
if (y % 2 == 1)
res = (res * x) % p;
y = y >> 1;
x = (x * x) % p;
}
return res;
}
public int WaysToBuildRooms(int[] prevRoom) {
int n = prevRoom.Length;
List<List<int>> tree = new List<List<int>>();
for (int i = 0; i < n; i++) tree.Add(new List<int>());
for (int i = 1; i < n; i++) tree[prevRoom[i]].Add(i);
return (int)DFS(tree, 0)[0];
}
private long[] DFS(List<List<int>> tree, int node) {
long way = 1;
int size = 0;
foreach (int ch in tree[node]) {
long[] sub = DFS(tree, ch);
way = way * Comb(size + (int)sub[1], (int)sub[1]) % MOD * sub[0] % MOD;
size += (int)sub[1];
}
return new long[]{way, size + 1};
}
}This Python solution defines the waysToBuildRooms method, which first sets up a tree based on the prevRoom relationships. It then performs a DFS traversal and calculates the number of ways to build the subtree from each node using combinatorics to account for different valid orderings of child nodes.
The time complexity is O(n log n) due to the combination calculations, and the space complexity is O(n) due to the recursion stack and storage of the tree structure.
Approach 4: Approach 2: Topological Sort with Dynamic Programming
This approach leverages topological sorting to handle the sequential dependencies defined by prevRoom. Using dynamic programming, it calculates the number of ways to validly sequence room constructions such that all dependencies are respected. We calculate the result by processing nodes in topological order, ensuring all dependencies are met iteratively, hence constructing the final permutations.
from collections import defaultdict, deque
MOD = 10**9 + 7
class Solution:
def factorial(self, n):
fact = [1] * (n + 1)
for i in range(2, n + 1):
fact[i] = fact[i - 1] * i % MOD
return fact
def mod_inverse(self, a, p):
return pow(a, p - 2, p)
def waysToBuildRooms(self, prevRoom):
n = len(prevRoom)
adj = defaultdict(list)
indegree = [0] * n
for i in range(1, n):
adj[prevRoom[i]].append(i)
indegree[i] += 1
zero_indegree = deque([i for i in range(n) if indegree[i] == 0])
order = []
while zero_indegree:
node = zero_indegree.popleft()
order.append(node)
for neighbor in adj[node]:
indegree[neighbor] -= 1
if indegree[neighbor] == 0:
zero_indegree.append(neighbor)
fact = self.factorial(n)
dp = [0] * n
dp[0] = 1
for node in order:
num_subnodes = 0
product = 1
for neighbor in adj[node]:
num_subnodes += dp[neighbor]
product = product * dp[neighbor] % MOD
dp[node] = product * fact[num_subnodes] % MOD * self.mod_inverse(fact[sum(dp[n] for n in adj[node])], MOD) % MOD
return dp[0]#include <iostream>
#include <vector>
#include <queue>
using namespace std;
static const int MOD = 1e9 + 7;
long long power(long long x, long long y, long long p) {
long long res = 1;
x = x % p;
while (y > 0) {
if (y % 2 == 1) res = (res * x) % p;
y = y >> 1;
x = (x * x) % p;
}
return res;
}
long long factorial(int n) {
long long fact = 1;
for (int i = 2; i <= n; i++) fact = fact * i % MOD;
return fact;
}
class Solution {
public:
int waysToBuildRooms(vector<int>& prevRoom) {
int n = prevRoom.size();
vector<vector<int>> adj(n);
vector<int> indegree(n, 0);
for (int i = 1; i < n; i++) {
adj[prevRoom[i]].push_back(i);
indegree[i]++;
}
queue<int> zeroIndegree;
for (int i = 0; i < n; i++) if (indegree[i] == 0) zeroIndegree.push(i);
vector<int> order;
while (!zeroIndegree.empty()) {
int node = zeroIndegree.front();
zeroIndegree.pop();
order.push_back(node);
for (int neighbor : adj[node]) {
indegree[neighbor]--;
if (indegree[neighbor] == 0) zeroIndegree.push(neighbor);
}
}
vector<long long> dp(n, 0);
dp[0] = 1;
vector<long long> fact(n + 1, 1);
for (int i = 2; i <= n; i++) fact[i] = fact[i - 1] * i % MOD;
for (int i = 0; i < order.size(); i++) {
int node = order[i];
long long numSubNodes = 0;
long long product = 1;
for (int child : adj[node]) {
product = product * dp[child] % MOD;
numSubNodes += dp[child];
}
dp[node] = product * fact[numSubNodes] % MOD * power(fact[numSubNodes], MOD - 2, MOD) % MOD;
}
return dp[0];
}
};In this Python solution, the tree is designed as an adjacency list based on prevRoom. It identifies all nodes with zero indegrees for topological processing. It then dynamically determines valid sequences, considering factorial-based permutations of child orderings, and adjusts based on dependencies, accumulating into dp[node].
The time complexity lies at O(n) due to the single traversal and computation using topological order, and the space complexity also holds at O(n) from storage of dependencies and intermediate states.
Complexity comparison
| Approach | Complexity |
|---|---|
| Dynamic Programming with Tree Representation | Time Complexity: O(n^2), where n is the number of rooms, due to combinations precomputation. Space Complexity: O(n^2) for the combination table to store computed values. |
| Combination Mathematics in Tree | Time Complexity: O(n^2), due to combination calculations. Space Complexity: O(n^2), used for the combination cache. |
| Approach 1: DFS and Combinatorics | The time complexity is O(n log n) due to the combination calculations, and the space complexity is O(n) due to the recursion stack and storage of the tree structure. |
| Approach 2: Topological Sort with Dynamic Programming | The time complexity lies at O(n) due to the single traversal and computation using topological order, and the space complexity also holds at O(n) from storage of dependencies and intermediate states. |
Complexity analysis
| Approach | Time | Space | When to Use |
|---|---|---|---|
| DFS and Combinatorics | O(n) | O(n) | Best general solution; combines subtree sizes with multinomial coefficients |
| Dynamic Programming on Tree | O(n) | O(n) | When implementing structured subtree DP with explicit states |
| Combination Mathematics in Tree | O(n) | O(n) | When focusing on mathematical formulation of valid interleavings |
| Topological Sort with DP | O(n) | O(n) | Useful for reasoning about dependency ordering in DAG-style problems |
Frequently asked questions
Is Count Ways to Build Rooms in an Ant Colony easy or hard?
Count Ways to Build Rooms in an Ant Colony is classified as Hard. The challenge comes from recognizing that valid build sequences correspond to interleavings of independent subtrees and implementing the combinatorics efficiently with modular arithmetic.
Count Ways to Build Rooms in an Ant Colony Python/Java solution
Most implementations use DFS on the tree combined with factorial-based combination calculations under modulo. Python, Java, C++, and JavaScript versions follow the same structure: build adjacency lists, precompute factorials, run DFS, and combine subtree counts using multinomial coefficients.
How to solve Count Ways to Build Rooms in an Ant Colony in O(n)?
Build a tree from the parent array and run DFS from the root. For each node, calculate the number of nodes in its subtree and multiply child results while using combinations to interleave their build orders. Using precomputed factorials and inverse factorials allows all combinatorial calculations in constant time.
What is the best approach for Count Ways to Build Rooms in an Ant Colony?
DFS with combinatorics on the tree is the most efficient approach. Compute subtree sizes using DFS and count ways to interleave children using multinomial coefficients. Precomputing factorials and inverse factorials allows combination calculations in O(1), giving an overall O(n) solution.
Is Count Ways to Build Rooms in an Ant Colony asked at Google/Amazon/Meta?
Tree dynamic programming and combinatorics problems like this commonly appear in interviews at companies such as Google, Meta, and Amazon. The problem tests understanding of dependency graphs, subtree aggregation, and modular combinatorics.
What data structure is used in Count Ways to Build Rooms in an Ant Colony?
The core data structure is a tree represented with an adjacency list built from the parent array. The algorithm also uses arrays for factorials, inverse factorials, and subtree sizes to compute combinations efficiently.
What is the time complexity of Count Ways to Build Rooms in an Ant Colony?
The optimal solution runs in O(n) time. Each node is visited once during DFS, and combination values are computed using precomputed factorials and modular inverses. Space complexity is O(n) for the adjacency list, factorial arrays, and recursion stack.