Use this to learn the idea, then write your own version.
123 45class Solution(object):6 def finishTime(self, n, edges, baseTime):7 """8 :type n: int9 :type edges: List[List[int]]10 :type baseTime: List[int]11 :rtype: int12 """13 POS_INF, NEG_INF = float("inf"), float("-inf")14 def iter_dfs():15 dp = [0]*n16 stk = [(1, 0, -1)]17 while stk:18 step, u, p = stk.pop()19 if step == 1:20 stk.append((2, u, p))21 for v in reversed(adj[u]):22 if v == p:23 continue24 stk.append((1, v, u))25 elif step == 2:26 mx, mn = NEG_INF, POS_INF27 for v in adj[u]:28 if v == p:29 continue30 mx, mn = max(mx, dp[v]), min(mn, dp[v])31 dp[u] = ((2*mx-mn) if mx is not NEG_INF else 0)+baseTime[u]32 return dp33 34 def iter_dfs2():35 def top2(a, b, x, cmp):36 if cmp(x, a):37 a, b = x, a38 elif cmp(x, b):39 b = x40 return a, b41 42 result = POS_INF43 stk = [(0, -1, NEG_INF)]44 while stk:45 u, p, t = stk.pop()46 mx1, mx2, mn1, mn2 = NEG_INF, NEG_INF, POS_INF, POS_INF47 for v in adj[u]:48 x = dp[v] if v != p else t49 mx1, mx2 = top2(mx1, mx2, x, lambda x, y: x > y)50 mn1, mn2 = top2(mn1, mn2, x, lambda x, y: x < y)51 result = min(result, ((2*mx1-mn1) if mx1 is not NEG_INF else 0)+baseTime[u])52 for v in reversed(adj[u]):53 if v == p:54 continue55 mx = mx1 if dp[v] != mx1 else mx256 mn = mn1 if dp[v] != mn1 else mn257 stk.append((v, u, ((2*mx-mn) if mx is not NEG_INF else 0)+baseTime[u]))58 return result59 60 adj = [[] for _ in xrange(n)]61 for u, v in edges:62 adj[u].append(v)63 adj[v].append(u)64 dp = iter_dfs()65 return iter_dfs2()66 67 68697071class Solution2(object):72 def finishTime(self, n, edges, baseTime):73 """74 :type n: int75 :type edges: List[List[int]]76 :type baseTime: List[int]77 :rtype: int78 """79 POS_INF, NEG_INF = float("inf"), float("-inf")80 def dfs(u, p):81 mx, mn = NEG_INF, POS_INF82 for v in adj[u]:83 if v == p:84 continue85 dfs(v, u)86 mx, mn = max(mx, dp[v]), min(mn, dp[v])87 dp[u] = ((2*mx-mn) if mx is not NEG_INF else 0)+baseTime[u]88 89 def dfs2(u, p, t):90 def top2(a, b, x, cmp):91 if cmp(x, a):92 a, b = x, a93 elif cmp(x, b):94 b = x95 return a, b96 97 mx1, mx2, mn1, mn2 = NEG_INF, NEG_INF, POS_INF, POS_INF98 for v in adj[u]:99 x = dp[v] if v != p else t100 mx1, mx2 = top2(mx1, mx2, x, lambda x, y: x > y)101 mn1, mn2 = top2(mn1, mn2, x, lambda x, y: x < y)102 result[0] = min(result[0], ((2*mx1-mn1) if mx1 is not NEG_INF else 0)+baseTime[u])103 for v in adj[u]:104 if v == p:105 continue106 mx = mx1 if dp[v] != mx1 else mx2107 mn = mn1 if dp[v] != mn1 else mn2108 dfs2(v, u, ((2*mx-mn) if mx is not NEG_INF else 0)+baseTime[u])109 110 adj = [[] for _ in xrange(n)]111 for u, v in edges:112 adj[u].append(v)113 adj[v].append(u)114 dp = [0]*n115 dfs(0, -1)116 result = [POS_INF]117 dfs2(0, -1, NEG_INF)118 return result[0]119