Create count-ways-to-build-rooms-in-an-ant-colony.py

Author: kamyu104

Python/count-ways-to-build-rooms-in-an-ant-colony.py

@@ -0,0 +1,32 @@
+# Time:  O(n)
+# Space: O(n)
+
+class Solution(object):
+    def waysToBuildRooms(self, prevRoom):
+        """
+        :type prevRoom: List[int]
+        :rtype: int
+        """
+        MOD = 10**9+7
+        fact = [1, 1]
+        inv = [0, 1]
+        inv_fact = [1, 1]
+        def nCr(n, k):
+            while len(inv) <= n:  # lazy initialization
+                fact.append(fact[-1]*len(inv) % MOD)
+                inv.append(inv[MOD%len(inv)]*(MOD-MOD//len(inv)) % MOD)  # https://cp-algorithms.com/algebra/module-inverse.html
+                inv_fact.append(inv_fact[-1]*inv[-1] % MOD)
+            return (fact[n]*inv_fact[n-k] % MOD) * inv_fact[k] % MOD
+
+        def dfs(adj, curr):
+            total_ways, total_cnt = 1, 0
+            for child in adj[curr]:
+                ways, cnt = dfs(adj, child)
+                total_cnt += cnt
+                total_ways = (((total_ways*ways) % MOD)*nCr(total_cnt, cnt)) % MOD
+            return total_ways, total_cnt+1
+
+        adj = [[] for _ in xrange(len(prevRoom))]
+        for i in xrange(1, len(prevRoom)):
+            adj[prevRoom[i]].append(i)
+        return dfs(adj, 0)[0]