1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46
| class Solution: def lowestCommonAncestor(self, root: 'TreeNode', p: 'TreeNode', q: 'TreeNode') -> 'TreeNode': ''' # Recursive def dfs(cur): if cur is None: return 0, None cnt_l, node_l = dfs(cur.left) cnt_r, node_r = dfs(cur.right) if cnt_l == 2: return 2, node_l if cnt_r == 2: return 2, node_r
if (cnt_l+cnt_r+(cur.val in [p.val, q.val])) == 2: return 2, cur return (1, cur) if cur.val in [p.val, q.val] else ((cnt_l, node_l) if cnt_l else (cnt_r, node_r)) return dfs(root)[1] ''' ret = [0] * 4 stack, cur = [], ((root, [0] * 4), ret, 'l') while cur[0][0] or stack: while cur[0][0]: stack.append(cur) if cur[0][0].left: cur = ((cur[0][0].left, [0]*4), cur[0][1], 'l') else: cur = ((cur[0][0].right, [0]*4), cur[0][1], 'r') (node, (cnt_l, node_l, cnt_r, node_r)), p_state, dir = stack.pop()
if cnt_l == 2: my_ret = 2, node_l elif cnt_r == 2: my_ret = 2, node_r elif (cnt_l+cnt_r+(node.val in [p.val, q.val])) == 2: my_ret = 2, node else: my_ret = (1, node) if node.val in [p.val, q.val] else ((cnt_l, node_l) if cnt_l else (cnt_r, node_r))
if dir == 'l': p_state[:2] = my_ret else: p_state[2:] = my_ret if stack and stack[-1][0][0].left == node: cur = ((stack[-1][0][0].right, [0] * 4), stack[-1][0][1], 'r') else: cur = ((None, None), (None, None), None)
return ret[1]
|