Count Ways to Build Rooms in an Ant Colony
HardUpdated: Aug 2, 2025
Practice on:
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 thenumber of different orders you can build all the rooms in.
Since the answer may be large, return it modulo 10^9 + 7.
Examples
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 <= 10^5prevRoom[0] == -10 <= prevRoom[i] < nfor all1 <= i < n- Every room is reachable from room
0once all the rooms are built.
Solution
Method 1 – Tree DP with Combinatorics (Topological Orderings)
Intuition
The problem reduces to counting the number of valid topological orderings of a rooted tree, where each room can only be built after its parent. The number of ways to build all rooms is the product of the ways to build each subtree, multiplied by the multinomial coefficient for merging the subtrees.
Approach
- Build the tree from the
prevRoomarray. - Precompute factorials and inverse factorials modulo for multinomial coefficients.
- Use DFS to compute for each node:
- The size of its subtree.
- The number of ways to build all rooms in its subtree.
- For each node, the answer is the product of the answers for its children, multiplied by the multinomial coefficient for merging the subtrees.
- Return the answer for the root node.
Code
C++
class Solution {
public:
int waysToBuildRooms(vector<int>& prevRoom) {
const int MOD = 1e9+7, n = prevRoom.size();
vector<vector<int>> g(n);
for (int i = 1; i < n; ++i) g[prevRoom[i]].push_back(i);
vector<long long> fact(n+1, 1), invFact(n+1, 1);
for (int i = 1; i <= n; ++i) fact[i] = fact[i-1] * i % MOD;
auto modinv = [&](long long x) {
long long res = 1, p = MOD-2;
while (p) { if (p&1) res = res*x%MOD; x = x*x%MOD; p>>=1; }
return res;
};
for (int i = 1; i <= n; ++i) invFact[i] = modinv(fact[i]);
function<pair<int,long long>(int)> dfs = [&](int u) {
int sz = 1;
long long ways = 1;
for (int v : g[u]) {
auto [csz, cways] = dfs(v);
ways = ways * cways % MOD * invFact[csz] % MOD;
sz += csz;
}
ways = ways * fact[sz-1] % MOD;
return make_pair(sz, ways);
};
return dfs(0).second;
}
};
Go
func waysToBuildRooms(prevRoom []int) int {
MOD := int(1e9 + 7)
n := len(prevRoom)
g := make([][]int, n)
for i := 1; i < n; i++ {
g[prevRoom[i]] = append(g[prevRoom[i]], i)
}
fact := make([]int64, n+1)
invFact := make([]int64, n+1)
fact[0] = 1
for i := 1; i <= n; i++ {
fact[i] = fact[i-1] * int64(i) % int64(MOD)
}
modinv := func(x int64) int64 {
res, p := int64(1), int64(MOD-2)
for p > 0 {
if p&1 == 1 {
res = res * x % int64(MOD)
}
x = x * x % int64(MOD)
p >>= 1
}
return res
}
for i := 0; i <= n; i++ {
invFact[i] = modinv(fact[i])
}
var dfs func(int) (int, int64)
dfs = func(u int) (int, int64) {
sz, ways := 1, int64(1)
for _, v := range g[u] {
csz, cways := dfs(v)
ways = ways * cways % int64(MOD) * invFact[csz] % int64(MOD)
sz += csz
}
ways = ways * fact[sz-1] % int64(MOD)
return sz, ways
}
_, ans := dfs(0)
return int(ans)
}
Java
class Solution {
public int waysToBuildRooms(int[] prevRoom) {
int MOD = 1_000_000_007, n = prevRoom.length;
List<Integer>[] g = new List[n];
for (int i = 0; i < n; i++) g[i] = new ArrayList<>();
for (int i = 1; i < n; i++) g[prevRoom[i]].add(i);
long[] fact = new long[n+1], invFact = new long[n+1];
fact[0] = 1;
for (int i = 1; i <= n; i++) fact[i] = fact[i-1] * i % MOD;
for (int i = 0; i <= n; i++) invFact[i] = modinv(fact[i], MOD);
return (int)dfs(0, g, fact, invFact, MOD)[1];
}
private long[] dfs(int u, List<Integer>[] g, long[] fact, long[] invFact, int MOD) {
int sz = 1;
long ways = 1;
for (int v : g[u]) {
long[] res = dfs(v, g, fact, invFact, MOD);
int csz = (int)res[0];
long cways = res[1];
ways = ways * cways % MOD * invFact[csz] % MOD;
sz += csz;
}
ways = ways * fact[sz-1] % MOD;
return new long[]{sz, ways};
}
private long modinv(long x, int MOD) {
long res = 1, p = MOD-2;
while (p > 0) {
if ((p&1) == 1) res = res * x % MOD;
x = x * x % MOD;
p >>= 1;
}
return res;
}
}
Kotlin
class Solution {
fun waysToBuildRooms(prevRoom: IntArray): Int {
val MOD = 1_000_000_007
val n = prevRoom.size
val g = Array(n) { mutableListOf<Int>() }
for (i in 1 until n) g[prevRoom[i]].add(i)
val fact = LongArray(n+1) { 1L }
val invFact = LongArray(n+1) { 1L }
for (i in 1..n) fact[i] = fact[i-1] * i % MOD
fun modinv(x: Long): Long {
var res = 1L; var p = MOD-2L; var b = x
while (p > 0) {
if (p and 1L == 1L) res = res * b % MOD
b = b * b % MOD
p = p shr 1
}
return res
}
for (i in 1..n) invFact[i] = modinv(fact[i])
fun dfs(u: Int): Pair<Int, Long> {
var sz = 1
var ways = 1L
for (v in g[u]) {
val (csz, cways) = dfs(v)
ways = ways * cways % MOD * invFact[csz] % MOD
sz += csz
}
ways = ways * fact[sz-1] % MOD
return sz to ways
}
return dfs(0).second.toInt()
}
}
Python
class Solution:
def waysToBuildRooms(self, prevRoom: list[int]) -> int:
MOD = 10**9 + 7
n = len(prevRoom)
g = [[] for _ in range(n)]
for i in range(1, n):
g[prevRoom[i]].append(i)
fact = [1] * (n+1)
invFact = [1] * (n+1)
for i in range(1, n+1):
fact[i] = fact[i-1] * i % MOD
def modinv(x):
res, p = 1, MOD-2
while p:
if p&1: res = res*x%MOD
x = x*x%MOD
p //= 2
return res
for i in range(1, n+1):
invFact[i] = modinv(fact[i])
def dfs(u):
sz, ways = 1, 1
for v in g[u]:
csz, cways = dfs(v)
ways = ways * cways % MOD * invFact[csz] % MOD
sz += csz
ways = ways * fact[sz-1] % MOD
return sz, ways
return dfs(0)[1]
Rust
impl Solution {
pub fn ways_to_build_rooms(prev_room: Vec<i32>) -> i32 {
const MOD: i64 = 1_000_000_007;
let n = prev_room.len();
let mut g = vec![vec![]; n];
for i in 1..n {
g[prev_room[i] as usize].push(i);
}
let mut fact = vec![1i64; n+1];
let mut inv_fact = vec![1i64; n+1];
for i in 1..=n {
fact[i] = fact[i-1] * i as i64 % MOD;
}
fn modinv(mut x: i64) -> i64 {
let mut res = 1i64; let mut p = MOD-2;
while p > 0 {
if p & 1 == 1 { res = res * x % MOD; }
x = x * x % MOD;
p >>= 1;
}
res
}
for i in 1..=n {
inv_fact[i] = modinv(fact[i]);
}
fn dfs(u: usize, g: &Vec<Vec<usize>>, fact: &Vec<i64>, inv_fact: &Vec<i64>) -> (usize, i64) {
let mut sz = 1;
let mut ways = 1i64;
for &v in &g[u] {
let (csz, cways) = dfs(v, g, fact, inv_fact);
ways = ways * cways % MOD * inv_fact[csz] % MOD;
sz += csz;
}
ways = ways * fact[sz-1] % MOD;
(sz, ways)
}
dfs(0, &g, &fact, &inv_fact).1 as i32
}
}
TypeScript
class Solution {
waysToBuildRooms(prevRoom: number[]): number {
const MOD = 1e9 + 7
const n = prevRoom.length
const g: number[][] = Array.from({length: n}, () => [])
for (let i = 1; i < n; ++i) g[prevRoom[i]].push(i)
const fact = Array(n+1).fill(1)
const invFact = Array(n+1).fill(1)
for (let i = 1; i <= n; ++i) fact[i] = fact[i-1] * i % MOD
function modinv(x: number): number {
let res = 1, p = MOD-2
while (p) {
if (p&1) res = res*x%MOD
x = x*x%MOD
p >>= 1
}
return res
}
for (let i = 1; i <= n; ++i) invFact[i] = modinv(fact[i])
function dfs(u: number): [number, number] {
let sz = 1, ways = 1
for (const v of g[u]) {
const [csz, cways] = dfs(v)
ways = ways * cways % MOD * invFact[csz] % MOD
sz += csz
}
ways = ways * fact[sz-1] % MOD
return [sz, ways]
}
return dfs(0)[1]
}
}
Complexity
- ⏰ Time complexity:
O(n)for tree traversal and factorial precomputation. - 🧺 Space complexity:
O(n)for tree, factorials, and recursion stack.