Update 1916.Count Ways to Build Rooms in an Ant Colony.py

Author: yanglbme

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

@@ -1,31 +1,26 @@
-from collections import defaultdict
-from math import comb
-from typing import List
-
-
 class Solution:
-  def waysToBuildRooms(self, prevRoom: List[int]) -> int:
-    modulo = 10 ** 9 + 7
-    ingoing = defaultdict(set)
-    outgoing = defaultdict(set)
+    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]
+        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  # just self
+        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
+            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
+        recurse(0)
+        return ans[0] % modulo