Commit 9dcc160

authored

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

1 parent 5ccbfdd commit 9dcc160Copy full SHA for 9dcc160

File tree

+31

-0

lines changed

+31

-0

lines changed

Lines changed: 31 additions & 0 deletions

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,31 @@`
	`1`	`+from collections import defaultdict`
	`2`	`+from math import comb`
	`3`	`+from typing import List`
	`4`	`+`
	`5`	`+`
	`6`	`+class Solution:`
	`7`	`+ def waysToBuildRooms(self, prevRoom: List[int]) -> int:`
	`8`	`+ modulo = 10 ** 9 + 7`
	`9`	`+ ingoing = defaultdict(set)`
	`10`	`+ outgoing = defaultdict(set)`
	`11`	`+`
	`12`	`+ for i in range(1, len(prevRoom)):`
	`13`	`+ ingoing[i].add(prevRoom[i])`
	`14`	`+ outgoing[prevRoom[i]].add(i)`
	`15`	`+ ans = [1]`
	`16`	`+`
	`17`	`+ def recurse(i):`
	`18`	`+ if len(outgoing[i]) == 0:`
	`19`	`+ return 1 # just self`
	`20`	`+`
	`21`	`+ nodes_in_tree = 0`
	`22`	`+ for v in outgoing[i]:`
	`23`	`+ cn = recurse(v)`
	`24`	`+ if nodes_in_tree != 0:`
	`25`	`+ ans[0] *= comb(nodes_in_tree + cn, cn)`
	`26`	`+ ans[0] %= modulo`
	`27`	`+ nodes_in_tree += cn`
	`28`	`+ return nodes_in_tree + 1`
	`29`	`+`
	`30`	`+ recurse(0)`
	`31`	`+ return ans[0] % modulo`

Comments

(0)