leetcode-bt

236. Lowest Common Ancestor of a Binary Tree.py

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

class Solution:
    """
    Thought Process:
        Post-order traversal since we want to find lowest ancestoer (root node last)
        
        For each node, we search if p or q in the left subtree and did the same thing in the right subtree
      1  ->left = 1, right = None
     /
    2  -> left=None, right=None
    and given 1 and 2 as p and q -> return 1

      3  ->left = None, right = 1
     / \
    5   1 -> left=None, right=None
-> left=None, right=None
    and given 1 and 3 as p and q -> return 3 (current root node)

      3  ->left = 5, right = 1
     / \
    5   1 -> left=None, right=None
-> left=None, right=None
    and given 5 and 1 as p and q -> return 3 (current root node)

    """
    def lowestCommonAncestor(self, root: 'TreeNode', p: 'TreeNode', q: 'TreeNode') -> 'TreeNode':
        if not root:
            return None
        def dfs(r):
            """post-order"""
            if not r:
                return None
            if r.val == p.val or r.val == q.val:
                return r
            left = dfs(r.left)
            right = dfs(r.right)
            if left and right:
                # if we both can find p in the one subtree and q in another subtree, found the ancestor which is root node
                return r
            if not left:
                # I didn not find anything in the left
                return right
            if not right:
                return left
            # it's impossible to not find anything in both left and right subtre. Otherwise, there's no ancestor
        return dfs(root)
    def lowestCommonAncestor(self, root: 'TreeNode', p: 'TreeNode', q: 'TreeNode') -> 'TreeNode':
        """
        For each current given root node, we try to find p in the left subtree and q in the right subtree.
        T:O(n). In worse case, each node would be visited.
        S:O(h) because of recursion call.
        """
        def dfs(r):
            if not r:
                return None
            if r.val == p.val or r.val == q.val:
                return r
            left = dfs(r.left)
            right = dfs(r.right)
            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
        return dfs(root)