Count ways to build rooms in an ant colony

Array Math Dynamic Programming Tree Depth-First Search Graph Theory Topological Sort Combinatorics

51.1% acceptance

Feb 25, 2026

524

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.

Solution

Rust

Time O(n * m)

Space O(n * m)

LeetCode

solution.rs

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 children = vec![vec![]; n];
    for i in 1..n {
      children[prev_room[i] as usize].push(i);
    }

    // Precompute factorials and inverse factorials
    let mut fact = vec![1i64; n + 1];
    for i in 1..=n {
      fact[i] = fact[i - 1] * i as i64 % MOD;
    }
    let mut inv_fact = vec![1i64; n + 1];
    inv_fact[n] = Self::power(fact[n], MOD - 2, MOD);
    for i in (0..n).rev() {
      inv_fact[i] = inv_fact[i + 1] * (i + 1) as i64 % MOD;
    }

    let mut size = vec![0usize; n];
    let mut ans = vec![1i64; n];

    // Iterative post-order traversal
    let mut stack: Vec<(usize, bool)> = vec![(0, false)];
    while let Some((node, processed)) = stack.pop() {
      if processed {
        size[node] = 1;
        let mut total = 0usize;
        for &child in &children[node] {
          total += size[child];
          ans[node] = ans[node] * ans[child] % MOD;
          ans[node] = ans[node] * inv_fact[size[child]] % MOD;
        }
        ans[node] = ans[node] * fact[total] % MOD;
        size[node] = total + 1;
      } else {
        stack.push((node, true));
        for &child in &children[node] {
          stack.push((child, false));
        }
      }
    }

    ans[0] as i32
  }

  fn power(mut base: i64, mut exp: i64, modulus: i64) -> i64 {
    let mut result = 1i64;
    base %= modulus;
    while exp > 0 {
      if exp & 1 == 1 {
        result = result * base % modulus;
      }
      exp >>= 1;
      base = base * base % modulus;
    }
    result
  }
}