1644. Lowest Common Ancestor of a Binary Tree II - cocoder39/coco39_LC GitHub Wiki

1644. Lowest Common Ancestor of a Binary Tree II

• first get potential LCA by leveraging regular LCA process
  • if p is returned, check if q is in subtree of p
  • if q is returned, check if p is in subtree of q
  • otherwise return the candidate lca
    • if lca is none, then no lca, then return none
    • if lca is not none, and it is not p or q, then it is lca

class Solution:
    def lowestCommonAncestor(self, root: 'TreeNode', p: 'TreeNode', q: 'TreeNode') -> 'TreeNode':
        def lca(root, p, q):
            if not root:
                return None
            
            if root == p or root == q:
                return root
            
            left = lca(root.left, p, q)
            right = lca(root.right, p, q)
            if left and right:
                return root
            if left:
                return left
            return right
        
        def contains(root, target):
            if not root:
                return False
            
            if root == target:
                return True
            return contains(root.left, target) or contains(root.right, target)
        
        lca = lca(root, p, q)
        if lca == p:
            return p if contains(p, q) else None
        if lca == q:
            return q if contains(q, p) else None
        return lca