Leetcode 0236. Lowest Common Ancestor of a Binary Tree

Problem statement

https://leetcode.com/problems/lowest-common-ancestor-of-a-binary-tree/

Solution

The idea is to traverse our tree and keep auxilary function dfs(node), which returns 0 if subtree do not have desired nodes and 1 if it has (one or two of them). For given node, we calculate information about:

1. lft is dfs(node.left), that is answer for left subtree.
2. rgh is dfs(node.right), that is answer for right subtree.
3. mid is true if current node is either p or q.

Then what we need to find is such node, for which at least 2 out of these 3 values are true. Why? If p is parent of q, then the answer is p and we have mid = 1 and one of the lft and rgh equal to 1. If q is parent of p, it is similar. If we have h which is LCA of p and q, then for h we have lft = rgh = 1 and mid = 0.

Complexity

Complexity is O(n), both time and space.

Code

class Solution:
    def lowestCommonAncestor(self, root, p, q):
        def dfs(node):
            if not node: return 0

            lft = dfs(node.left)
            rgh = dfs(node.right)
            mid = node in [p, q]
            if lft + rgh + mid >= 2:
                self.Found = node

            return max(lft, mid, rgh)
        
        self.Found = None
        dfs(root)
        return self.Found