Count Ways to Build Rooms in an Ant Colony

Count Ways to Build Rooms in an Ant Colony - Problem

Math Dynamic Programming Tree Graph Topological Sort Combinatorics Hard

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 10⁹ + 7.

Input & Output

Example 1 — Basic Tree

$ Input: prevRoom = [-1,0,1,2]

› Output: 1

💡 Note: Only one valid order: build rooms in sequence 0→1→2→3 since each room depends on the previous one

Example 2 — Multiple Children

$ Input: prevRoom = [-1,0,0,1,2]

› Output: 8

💡 Note: Room 0 has children 1,2. Room 1 has children 3,4. Multiple valid orders exist like: 0→1→2→3→4, 0→1→3→2→4, etc.

Example 3 — Single Path

$ Input: prevRoom = [-1,0]

› Output: 1

💡 Note: Only two rooms: 0 is already built, then build room 1. Only one possible order.

Constraints

1 ≤ n ≤ 10⁵
prevRoom[0] == -1
0 ≤ prevRoom[i] ≤ n - 1 for all 1 ≤ i ≤ n - 1
prevRoom represents a valid tree

Visualization

Tap to expand

Understanding the Visualization

Input Dependencies

prevRoom array defines parent-child relationships

Tree Structure

Convert to tree where each room depends on its parent

Count Arrangements

Calculate valid building orders respecting dependencies

Key Takeaway

🎯 Key Insight: Dependencies form a tree structure, and arrangements can be calculated using multinomial coefficients

Asked in

G Google 25 M Microsoft 18 a Amazon 15 F Facebook 12

The key insight is to recognize this as a tree combinatorics problem. Convert the dependency array into a tree structure, then use multinomial coefficients to count arrangements. For each node with subtrees of sizes [s1, s2, ..., sk], the number of ways is (s1+s2+...+sk)!/(s1!×s2!×...×sk!). Best approach uses DFS with mathematical formula. Time: O(n), Space: O(n)

Common Approaches

Approach	Time	Space	Notes
✓ Mathematical Combinatorics Formula	O(n)	O(n)	Direct calculation using multinomial coefficients for tree arrangements
DFS with Tree Structure	O(n²)	O(n)	Build tree from dependencies and use DFS to calculate arrangements
Brute Force - Generate All Permutations	O(n!)	O(n)	Generate all possible building orders and count valid ones

Mathematical Combinatorics Formula — Algorithm Steps

Build tree from dependencies
For each node, calculate subtree sizes
Apply multinomial formula: n!/(n1!×n2!×...×nk!)

Visualization

Tap to expand

Step-by-Step Walkthrough

Tree Structure

Identify parent-child relationships

Subtree Sizes

Calculate size of each subtree

Apply Formula

Use multinomial coefficient for arrangements

Code -

solution.c — C

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define MOD 1000000007
#define MAX_N 1000

long long modPow(long long base, long long exp, long long mod) {
    long long result = 1;
    while (exp > 0) {
        if (exp % 2 == 1) {
            result = (result * base) % mod;
        }
        base = (base * base) % mod;
        exp /= 2;
    }
    return result;
}

long long modInverse(long long a) {
    return modPow(a, MOD - 2, MOD);
}

int children[MAX_N][MAX_N];
int childCount[MAX_N];
long long fact[MAX_N + 1];

typedef struct {
    long long ways;
    int size;
} Result;

Result dfs(int node) {
    if (childCount[node] == 0) {
        Result r = {1, 1};
        return r;
    }
    
    long long result = 1;
    int subtreeSize = 0;
    int childSizes[MAX_N];
    int numChildren = 0;
    
    for (int i = 0; i < childCount[node]; i++) {
        Result childResult = dfs(children[node][i]);
        result = (result * childResult.ways) % MOD;
        childSizes[numChildren++] = childResult.size;
        subtreeSize += childResult.size;
    }
    
    if (subtreeSize > 0) {
        long long multinomial = fact[subtreeSize];
        for (int i = 0; i < numChildren; i++) {
            multinomial = (multinomial * modInverse(fact[childSizes[i]])) % MOD;
        }
        result = (result * multinomial) % MOD;
    }
    
    Result r = {result, subtreeSize + 1};
    return r;
}

int solution(int* prevRoom, int n) {
    if (n == 1) return 1;
    
    memset(childCount, 0, sizeof(childCount));
    
    for (int i = 1; i < n; i++) {
        children[prevRoom[i]][childCount[prevRoom[i]]++] = i;
    }
    
    fact[0] = 1;
    for (int i = 1; i <= n; i++) {
        fact[i] = (fact[i-1] * i) % MOD;
    }
    
    Result result = dfs(0);
    return result.ways;
}

int main() {
    char line[10000];
    fgets(line, sizeof(line), stdin);
    
    int prevRoom[MAX_N];
    int n = 0;
    
    char* token = strtok(line, "[],\n ");
    while (token != NULL) {
        prevRoom[n++] = atoi(token);
        token = strtok(NULL, "[],\n ");
    }
    
    printf("%d\n", solution(prevRoom, n));
    return 0;
}

Time & Space Complexity

Time Complexity

⏱️

O(n)

Single DFS traversal with constant work per node

✓ Linear Growth

Space Complexity

O(n)

Recursion stack and adjacency list storage

⚡ Linearithmic Space

28.0K Views

Medium Frequency

~35 min Avg. Time

892 Likes

Ln 1, Col 1

Smart Actions

💡 Explanation

AI Ready

💡 Suggestion Tab to accept Esc to dismiss

// Output will appear here after running code

Code Editor Closed

Click the red button to reopen

Input & Output

Constraints

Visualization

Related Problems

Common Approaches

Mathematical Combinatorics Formula — Algorithm Steps

Visualization

Code -

Time & Space Complexity

Select Compiler