Use this to learn the idea, then write your own version.
123 4import collections5 6 7class Solution(object):8 def minTime(self, n, edges, hasApple):9 """10 :type n: int11 :type edges: List[List[int]]12 :type hasApple: List[bool]13 :rtype: int14 """15 graph = collections.defaultdict(list)16 for u, v in edges:17 graph[u].append(v)18 graph[v].append(u)19 20 result = [0, 0]21 s = [(1, (-1, 0, result))]22 while s:23 step, params = s.pop()24 if step == 1:25 par, node, ret = params26 ret[:] = [0, int(hasApple[node])]27 for nei in reversed(graph[node]):28 if nei == par:29 continue30 new_ret = [0, 0]31 s.append((2, (new_ret, ret)))32 s.append((1, (node, nei, new_ret)))33 else:34 new_ret, ret = params35 ret[0] += new_ret[0]+new_ret[1]36 ret[1] |= bool(new_ret[0]+new_ret[1])37 return 2*result[0]38 39 404142class Solution_Recu(object):43 def minTime(self, n, edges, hasApple):44 """45 :type n: int46 :type edges: List[List[int]]47 :type hasApple: List[bool]48 :rtype: int49 """50 def dfs(graph, par, node, hasApple):51 result, extra = 0, int(hasApple[node])52 for nei in graph[node]:53 if nei == par:54 continue55 count, found = dfs(graph, node, nei, hasApple)56 result += count+found57 extra |= bool(count+found)58 return result, extra59 60 graph = collections.defaultdict(list)61 for u, v in edges:62 graph[u].append(v)63 graph[v].append(u)64 return 2*dfs(graph, -1, 0, hasApple)[0]65 66 676869class Solution2(object):70 def minTime(self, n, edges, hasApple):71 """72 :type n: int73 :type edges: List[List[int]]74 :type hasApple: List[bool]75 :rtype: int76 """77 graph = collections.defaultdict(list)78 for u, v in edges:79 graph[u].append(v)80 graph[v].append(u)81 82 result = [0]83 s = [(1, (-1, 0, result))]84 while s:85 step, params = s.pop()86 if step == 1:87 par, node, ret = params88 tmp = [int(hasApple[node])]89 s.append((3, (tmp, ret)))90 for nei in reversed(graph[node]):91 if nei == par:92 continue93 new_ret = [0]94 s.append((2, (new_ret, tmp, ret)))95 s.append((1, (node, nei, new_ret)))96 elif step == 2:97 new_ret, tmp, ret = params98 ret[0] += new_ret[0]99 tmp[0] |= bool(new_ret[0])100 else:101 tmp, ret = params102 ret[0] += tmp[0]103 return 2*max(result[0]-1, 0)104 105 106107108class Solution2_Recu(object):109 def minTime(self, n, edges, hasApple):110 """111 :type n: int112 :type edges: List[List[int]]113 :type hasApple: List[bool]114 :rtype: int115 """116 def dfs(graph, par, node, has_subtree):117 result, extra = 0, int(hasApple[node])118 for nei in graph[node]:119 if nei == par:120 continue121 count = dfs(graph, node, nei, hasApple)122 result += count123 extra |= bool(count)124 return result+extra125 126 graph = collections.defaultdict(list)127 for u, v in edges:128 graph[u].append(v)129 graph[v].append(u)130 return 2*max(dfs(graph, -1, 0, hasApple)-1, 0)131