Count Ways to Build Rooms in an Ant Colony LeetCode solution in Java

1
class Solution {
2
  int M = (int) 1e9 + 7;
3

4
  public int waysToBuildRooms(int[] prevRoom) {
5
    int n = prevRoom.length;
6
    long[] fact = new long[n];
7
    long[] invFact = new long[n];
8
    long[] inv = new long[n];
9
    fact[1] = fact[0] = invFact[0] = invFact[1] = inv[1] = 1;
10
    for (int i = 2; i < n; i++) { // modInverse
11
      fact[i] = fact[i - 1] * i % M;
12
      inv[i] = M - M / i * inv[M % i] % M;
13
      invFact[i] = invFact[i - 1] * inv[i] % M;
14
    }
15

16
    Map<Integer, List<Integer>> map = new HashMap<>();
17
    for (int i = 0; i < n; i++) { // add an edge from parent to child
18
      map.computeIfAbsent(prevRoom[i], o -> new ArrayList<>()).add(i);
19
    }
20

21
    long[] ans = new long[] {1};
22
    solve(0, fact, invFact, map, ans);
23
    return (int) ans[0];
24
  }
25

26
  private int solve(
27
      int i, long[] fact, long[] invFact, Map<Integer, List<Integer>> map, long[] ans) {
28
    int sum = 0;
29
    for (int next : map.getOrDefault(i, List.of())) {
30
      int cur = solve(next, fact, invFact, map, ans);
31
      ans[0] = ans[0] * invFact[cur] % M; // divide fact[cur] -> multiply invFact[cur]
32
      sum += cur;
33
    }
34
    ans[0] = ans[0] * fact[sum] % M;
35
    return sum + 1;
36
  }
37
}