Create Solution.py

Author: yanglbme

solution/1900-1999/1916.Count Ways to Build Rooms in an Ant Colony/Solution.py

@@ -0,0 +1,26 @@
+class Solution:
+    def waysToBuildRooms(self, prevRoom: List[int]) -> int:
+        modulo = 10**9 + 7
+        ingoing = defaultdict(set)
+        outgoing = defaultdict(set)
+
+        for i in range(1, len(prevRoom)):
+            ingoing[i].add(prevRoom[i])
+            outgoing[prevRoom[i]].add(i)
+        ans = [1]
+
+        def recurse(i):
+            if len(outgoing[i]) == 0:
+                return 1
+
+            nodes_in_tree = 0
+            for v in outgoing[i]:
+                cn = recurse(v)
+                if nodes_in_tree != 0:
+                    ans[0] *= comb(nodes_in_tree + cn, cn)
+                    ans[0] %= modulo
+                nodes_in_tree += cn
+            return nodes_in_tree + 1
+
+        recurse(0)
+        return ans[0] % modulo