Problem statement

https://leetcode.com/problems/count-ways-to-build-rooms-in-an-ant-colony/

Solution

The question is asked in this problem: given tree, how many ways there is to traverse it, so we traverse children after it parent. Imagine example: we have tree with root with 3 children, such that we have 10 nodes in total and 2, 4, 3 nodes for left, middle and right children. Then how many ways there is to traverse the tree? Let dp(i) be number of answers for subtree with node i. Then for the specific case we considered we can say that dp(0) = dp(1)*dp(2)*dp(3)*C_9^{2,4,3}, where C_9^{2,4,3} is multinomial coefficient = 9!/2!/4!/3!. Why it is so? Because we have 9 possible places to put 2 nodes of first type (left), 4 nodes of the second type (middle) and 3 nodes of the third type (right). Also we need to calculate number of solutions for smaller subproblems.

Here function dp(node) will return two values: first one is the size of subtree with this node and the second one is number of possible ways to built colony with root in this node. I precalculated all factorials and inverse factorials outside function: so in fact it is done only once. Here I use trick to evaluate inverses of all numbers from 1 to M in linear time, but actually we do not need this, we just need inverses of factorials which can be done simpler: evaluate inverse of M! and then multiply it to get inverse of (M-1)! and so on.

Complexity

Time complexity O(M) to calculate all factorials and inverses but we do it only once for all tests. Then we have O(n) to traverse our tree. Space complexity is O(M + n).

Code

class Solution:
    MOD, M = 10**9 + 7, 10**5
    inverses, inverse_fact, facts = [1]*(M+1), [1]*(M+1), [1]*(M+1)
    for a in range(2, M + 1):
        inverses[a] = MOD - ((MOD // a) * inverses[MOD % a]) % MOD

    for i in range(1, M+1):
        inverse_fact[i] = (inverse_fact[i-1] * inverses[i]) % MOD
        facts[i] = (facts[i-1] * i) % MOD

    def waysToBuildRooms(self, prevRoom):
        graph = defaultdict(list)
        for i, num in enumerate(prevRoom):
            graph[num].append(i)

        @lru_cache(None)
        def dp(node):
            ans, sm = 1, 1
            for neib in graph[node]:
                ans2, sm2 = dp(neib)
                ans = (ans * self.inverse_fact[sm2] * ans2) % self.MOD
                sm += sm2
            return ((ans * self.facts[sm-1]) % self.MOD, sm)

        return dp(0)[0]