LeetCode 1916. Count Ways to Build Rooms in an Ant Colony

Posted on Edited on In 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
class Solution:
    def waysToBuildRooms(self, prevRoom: List[int]) -> int:
        MOD = 10**9 + 7
        g = collections.defaultdict(list)
        for curr, prev in enumerate(prevRoom): g[prev].append(curr)
    
        def dfs(curr):
            if not g[curr]: return 1, 1
            ans, l = 1, 0
            for next in g[curr]:
                tmp, r = dfs(next)
                ans = (ans * tmp * math.comb(l + r, r)) % MOD
                l += r
            return ans, l + 1

        return dfs(0)[0]

Reference: [Python/C++] clean DFS solution with explanation