1644. Lowest Common Ancestor of a Binary Tree II.py

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

class Solution:
    """
    For each current given root node, we try to find p in the left subtree and q in the right subtree.
    However, since this problem does not guarantee that p and q must exist in the tree.
    So, we need two flags to keep track of if p or q exist in the tree.
    """
    def lowestCommonAncestor(self, root: 'TreeNode', p: 'TreeNode', q: 'TreeNode') -> 'TreeNode':
        self.flag_p = False
        self.flag_q = False
        def dfs(r):
            if not r:
                return None
            left = dfs(r.left)
            right = dfs(r.right)

            if r.val == p.val:
                self.flag_p = True
                return r
            if r.val == q.val:
                self.flag_q = True
                return r
            # find common ancestor
            if left and right:
                # r is common ancesor, if we found p in the left subtree and found q in the right subtree
                return r
            
            if not left:
                return right
            if not right:
                return left
        r = dfs(root)
        if not self.flag_p or not self.flag_q:
            return None
        return r