lca_basic.py

from dataclasses import dataclass

@dataclass
class TraversalRecord:
  seen_nodes: int
  ancestor: BinaryTreeNode


def post_order_traversal_helper(node: BinaryTreeNode, n1: BinaryTreeNode, n2: BinaryTreeNode) -> TraversalRecord:
  '''
  A helper method that constructs our TraversalRecords and checks that the binary tree
  '''
  if node is None:
    return TraversalRecord(0, None)
  
  left_traversal = post_order_traversal_helper(node.left)
  if left_traversal.ancestor is not None:
    return left_traversal

  right_traversal = post_order_traversal_helper(node.right)
  if right_traversal.ancestor is not None:
    return right_traversal
  
  # check if our current node is one of our target nodes
  seen_nodes = (n1, n2).count(node) \
    + left_traversal.seen_nodes + right_traversal.seen_nodes
  
  return TraversalRecord(
    seen_nodes=seen_nodes, 
    ancestor=node if seen_nodes == 2 else None
  )


def find_lca(root: BinaryTreeNode, n1: BinaryTreeNode, n2: BinaryTreeNode) -> BinaryTreeNode:
  return post_order_traversal_helper(root, n1, n2).ancestor