Firstly process the edge list and save as a node dict. Then start from node 0, DFS the whole tree. If current node is not root node and (current node has an apple or children nodes have apples), then the resultant time need to +2 (from parent node to current node and back).
1 #1443
2 #Runtime: 541 ms (Beats 100%)
3 #Memory: 54.5 MB (Beats 61.11%)
4
5 class Solution(object):
6 def minTime(self, n, edges, hasApple):
7 """
8 :type n: int
9 :type edges: List[List[int]]
10 :type hasApple: List[bool]
11 :rtype: int
12 """
13 nodes = defaultdict(list)
14 for x, y in edges:
15 nodes[x].append(y)
16 nodes[y].append(x)
17
18 def DFS(r, p):
19 res = 0
20 for son in nodes[r]:
21 if son != p:
22 res += DFS(son, r)
23 if r and (res or hasApple[r]):
24 return res + 2
25 return res
26
27 return DFS(0, -1)