Code Implementation

MOD = 10 ** 9 + 7

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)

        fact = [1] * (n + 1)
        inv_fact = [1] * (n + 1)
        for i in range(1, n + 1):
            fact[i] = fact[i - 1] * i % MOD
        inv_fact[n] = pow(fact[n], MOD - 2, MOD)
        for i in range(n, 0, -1):
            inv_fact[i - 1] = inv_fact[i] * i % MOD

        def dfs(u):
            total = 1
            size = 0
            for v in tree[u]:
                ways, sz = dfs(v)
                total = total * ways % MOD * inv_fact[sz] % MOD
                size += sz
            total = total * fact[size] % MOD
            return total, size + 1

        return dfs(0)[0]
    

#define MOD 1000000007
class Solution {
public:
    vector fact, inv_fact;
    int n;

    long long modpow(long long a, long long b) {
        long long res = 1;
        while (b) {
            if (b & 1) res = res * a % MOD;
            a = a * a % MOD;
            b >>= 1;
        }
        return res;
    }

    pair dfs(int u) {
        long long total = 1;
        int size = 0;
        for (int v : tree[u]) {
            auto [ways, sz] = dfs(v);
            total = total * ways % MOD * inv_fact[sz] % MOD;
            size += sz;
        }
        total = total * fact[size] % MOD;
        return {total, size + 1};
    }

    int waysToBuildRooms(vector<int>& prevRoom) {
        n = prevRoom.size();
        tree.assign(n, vector<int>());
        for (int i = 1; i < n; ++i) {
            tree[prevRoom[i]].push_back(i);
        }
        fact.assign(n + 1, 1);
        inv_fact.assign(n + 1, 1);
        for (int i = 1; i <= n; ++i)
            fact[i] = fact[i - 1] * i % MOD;
        inv_fact[n] = modpow(fact[n], MOD - 2);
        for (int i = n; i > 0; --i)
            inv_fact[i - 1] = inv_fact[i] * i % MOD;
        return dfs(0).first;
    }
};
    

class Solution {
    static final int MOD = 1_000_000_007;
    List<List<Integer>> tree;
    long[] fact, invFact;

    public int waysToBuildRooms(int[] prevRoom) {
        int n = prevRoom.length;
        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);

        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;
        invFact[n] = pow(fact[n], MOD - 2);
        for (int i = n; i > 0; i--) invFact[i - 1] = invFact[i] * i % MOD;

        return (int) dfs(0)[0];
    }

    private long[] dfs(int u) {
        long total = 1;
        int size = 0;
        for (int v : tree.get(u)) {
            long[] res = dfs(v);
            total = total * res[0] % MOD * invFact[(int) res[1]] % MOD;
            size += res[1];
        }
        total = total * fact[size] % MOD;
        return new long[]{total, size + 1};
    }

    private long pow(long a, long b) {
        long res = 1;
        while (b > 0) {
            if ((b & 1) == 1) res = res * a % MOD;
            a = a * a % MOD;
            b >>= 1;
        }
        return res;
    }
}
    

const MOD = 1e9 + 7;

var waysToBuildRooms = function(prevRoom) {
    const n = prevRoom.length;
    const tree = Array.from({length: n}, () => []);
    for (let i = 1; i < n; ++i) {
        tree[prevRoom[i]].push(i);
    }
    let fact = Array(n + 1).fill(1);
    let invFact = Array(n + 1).fill(1);
    for (let i = 1; i <= n; ++i)
        fact[i] = fact[i - 1] * i % MOD;
    invFact[n] = pow(fact[n], MOD - 2);
    for (let i = n; i > 0; --i)
        invFact[i - 1] = invFact[i] * i % MOD;

    function pow(a, b) {
        let res = 1;
        a %= MOD;
        while (b > 0) {
            if (b & 1) res = res * a % MOD;
            a = a * a % MOD;
            b >>= 1;
        }
        return res;
    }

    function dfs(u) {
        let total = 1, size = 0;
        for (let v of tree[u]) {
            let [ways, sz] = dfs(v);
            total = total * ways % MOD * invFact[sz] % MOD;
            size += sz;
        }
        total = total * fact[size] % MOD;
        return [total, size + 1];
    }

    return dfs(0)[0];
};
    

Problem Description

You are given an array prevRoom of length n, where prevRoom[i] indicates the parent room for room i in an ant colony. Room 0 is the root and has no parent (so prevRoom[0] = -1). Each new room can only be built after its parent room is built.

Your task is to count the total number of different valid orders in which all rooms can be built, following the parent-before-child constraint. The answer should be returned modulo 10^9 + 7.

• Each room can only be built once.
• You must respect the parent-child dependencies given by prevRoom.
• The input forms a tree rooted at room 0.

Thought Process

At first glance, the problem seems related to finding all possible topological sorts of a tree, since each room can only be built after its parent. A brute-force solution would be to generate all such orderings and count them, but this would be far too slow for large n (up to 10⁵).

A key insight is that the dependencies form a tree, and the orderings are determined by the ways we can interleave the construction of subtrees while always building parents before children. This is similar to counting linear extensions of a tree or the number of valid sequences respecting partial orders.

We need to count, for each subtree, the number of ways to build its rooms (recursively), and then combine the results for all subtrees of the current node, considering all the ways their construction orders can be interleaved while respecting dependencies.

Solution Approach

Let's break down the optimized approach step by step:

1. Build the tree: Create a tree structure from the prevRoom array, where each node points to its children.
2. Precompute factorials and inverse factorials: Since we need to count the number of ways to interleave sequences, we'll use combinations, which require factorials. We also need modular inverses for division under modulus.
3. DFS Traversal: Use depth-first search (DFS) to process each node from the leaves up to the root.
   • For each node, recursively compute the number of ways to build all its subtrees and the size (number of rooms) in each subtree.
   • To combine the subtrees, the number of ways to interleave their construction is given by the multinomial coefficient: fact[total size] / (fact[size1] * fact[size2] * ...).
   • Multiply the number of ways for each subtree with the multinomial coefficient to get the total for the current node.
4. Return the result: The answer is the number of ways to build all rooms starting from the root node.

We use arrays for factorials and their modular inverses for efficient computation. DFS ensures we process subproblems before combining them.

Example Walkthrough

Let's use a small example: prevRoom = [-1,0,1,2]

• Room 0 is the root.
• Room 1's parent is 0, room 2's parent is 1, room 3's parent is 2.

The tree is a straight line: 0 → 1 → 2 → 3

1. We must build room 0 first, then 1, then 2, then 3.
2. There is only one valid sequence: [0, 1, 2, 3]
3. The code will compute, for each node, that there is only one way to build its subtree (since there are no siblings to interleave).

Now consider prevRoom = [-1,0,0,1,1]:

• Room 0 is root, rooms 1 and 2 are children of 0, rooms 3 and 4 are children of 1.

The tree:
0
├── 1
│ ├── 3
│ └── 4
└── 2

1. We can build 1 or 2 after 0, but 3 and 4 only after 1.
2. For the subtree rooted at 1 (with 3 and 4), there are 2 ways (build 3 then 4, or 4 then 3).
3. We can interleave the sequence for subtree 1 (size 3) and subtree 2 (size 1) in C(3+1, 3) = 4 ways.
4. So, total = ways(1-subtree) * ways(2-subtree) * C(3+1, 3) = 2 * 1 * 4 = 8 ways.

Time and Space Complexity

Brute-force approach: Would attempt to generate all valid build orders, which is factorial in the number of nodes and not feasible for large n.

Optimized approach:

• Time Complexity: O(n), as each node is processed once in the DFS, and all factorials are precomputed in O(n).
• Space Complexity: O(n), for the tree structure, factorial arrays, and recursion stack.

Summary

The problem asks for the number of valid build orders for rooms in an ant colony, given parent-child dependencies. The key insight is that the number of valid orders is the product of the ways to build each subtree, multiplied by the number of ways to interleave their build sequences. By using DFS, multinomial coefficients, and precomputed factorials with modular inverses, we efficiently count all valid orders in linear time. This approach leverages combinatorial reasoning and tree traversal to solve a problem that would otherwise be intractable with brute force.