Welcome to Subscribe On Youtube
1916. Count Ways to Build Rooms in an Ant Colony
Description
You are an ant tasked with adding n
new rooms numbered 0
to n1
to your colony. You are given the expansion plan as a 0indexed 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 10^{9} + 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.length
2 <= n <= 10^{5}
prevRoom[0] == 1
0 <= prevRoom[i] < n
for all1 <= i < n
 Every room is reachable from room
0
once all the rooms are built.
Solution
Use dynamic programming based on topological sort. First, calculate each node’s indegree. Next, do topological sort starting on the nodes with 0 indegree, and do dynamic programming at the same time.
For dynamic programming, let dp[i]
represent the number of ways to build rooms with root at node i
. Initially, dp[i] = 1
for all 0 <= i < n
. For the nodes that are not leaves, calculate the dp
values according to combination values.
Finally, return dp[0]
.

class Solution { static final int MODULO = 1000000007; public int waysToBuildRooms(int[] prevRoom) { int n = prevRoom.length; int[] degree = new int[n]; for (int i = 0; i < n; i++) { if (prevRoom[i] >= 0) degree[prevRoom[i]]++; } int[] fac = new int[n + 1]; Arrays.fill(fac, 1); int[] inv = new int[n + 1]; Arrays.fill(inv, 1); for (int i = 2; i <= n; i++) fac[i] = (int) ((long) i * fac[i  1] % MODULO); for (int i = 2; i <= n; i++) inv[i] = (int) power(fac[i], MODULO  2); Queue<Integer> queue = new LinkedList<Integer>(); int[] dp = new int[n]; Arrays.fill(dp, 1); for (int i = 0; i < n; i++) { if (degree[i] == 0) queue.offer(i); } int[] sizes = new int[n]; Arrays.fill(sizes, 1); while (!queue.isEmpty()) { int u = queue.poll(); dp[u] = (int) ((long) dp[u] * fac[sizes[u]  1] % MODULO); int v = prevRoom[u]; if (v < 0) continue; sizes[v] += sizes[u]; if (degree[v] == 0) queue.offer(v); dp[v] = (int) ((long) dp[v] * inv[sizes[u]] % MODULO * dp[u] % MODULO); } return dp[0]; } public long power(long x, int n) { long pow = 1; for (int i = n; i != 0; i /= 2) { if (i % 2 == 1) pow = pow * x % MODULO; x = x * x % MODULO; } return pow; } }