leetcode-tree

1676. Lowest Common Ancestor of a Binary Tree IV.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 if given nodes in the left subtree and give nodes 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', nodes: 'List[TreeNode]') -> 'TreeNode':
        def dfs(r):
            if not r:
                return
            if r in nodes:
                return r
            
            left = dfs(r.left)
            right = dfs(r.right)
            if left and right:
                return r
            if not left:
                return right
            if not right:
                return left
        nodes = set(nodes)
        return dfs(root)