Algorithm Post: Divide and Conquer - Lowest Common Ancestor

lowest_common_ancestor.py

class TreeNode:
    def __init__(self, val):
        self.val = val
        self.left = self.right = None
    def __str__(self):
        return '[' + str(self.val) + ']'

class LCA:
    def find_lca(self, root, A, B):
        if not root:
            return
        if root is A or root is B:
            return root
        
        left = self.find_lca(root.left, A, B)
        right = self.find_lca(root.right, A, B)
        if left and right:
            return root
        return left or right
    
if __name__ == '__main__':
    trees = map(TreeNode, range(1, 6))
    trees[0].left = trees[1]
    trees[0].right = trees[2]
    trees[1].left = trees[3]
    trees[1].right = trees[4]
    
    print(LCA().find_lca(trees[0], trees[3], trees[4]))