236. Lowest Common Ancestor of a Binary Tree

# TC: O(N)
# SC: O(N), the height of the tree is how many recursion call stacks we incur. Worst case is N

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution:
    def lowestCommonAncestor(self, root: 'TreeNode', p: 'TreeNode', q: 'TreeNode') -> 'TreeNode':
        res = None

        def dfs(root, p, q) -> bool:
            nonlocal res
            if not root:
                return False

            found_root = root == p or root == q
            found_left = dfs(root.left, p, q)
            found_right = dfs(root.right, p, q)

            if found_left or found_right:
                if (found_left and found_right) or found_root:
                    res = root

            return found_root or found_left or found_right

        dfs(root, p, q)

        return res