1644. Lowest Common Ancestor of a Binary Tree II

# 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':
        """
        main idea: since p,q not necessary exist in the tree, need to traverse all the nodes in the tree before check if current node is p or q
        tc O(N) sc O(N)
        """
        self.cnt = 0
        def dfs(node,p,q):
            if not node: 
                return node
            # NOTE this is to find p or q in sub tree not finding LCA 
            leftpq = dfs(node.left,p,q)
            rightpq = dfs(node.right,p,q)
            # note the position of q,p check has to be below recursion but above subtree check 
            if p == node or q == node:
                self.cnt += 1
                return node 
            if leftLCA and rightLCA:
                return node
            elif rightLCA:
                return rightLCA
            elif leftLCA:
                return leftLCA

        res = dfs(root,p,q)
        return None if self.cnt != 2 else res