LC 1740 [M] Find Distance in a Binary Tree - ALawliet/algorithms GitHub Wiki

This code reuses the recursive function from Lowest Common Ancestor of Binary Tree (#236)

We want to find the LCA since it is guaranteed to be on the path between p and q. Since we're unable to move "up" in a binary tree, we can simple move down in both directions from the LCA to form the path. Return the distance from the LCA -> p, plus the distance from the LCA -> q.

Time complexity is O(n) since we have to look at all the nodes at least once in the worst case. Space complexity is O(logn) for the call stack.

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def findDistance(self, root: Optional[TreeNode], p: int, q: int) -> int:
        def postorder(node):
            if not node or node.val == p or node.val == q:
                return node
            L = postorder(node.left)
            R = postorder(node.right)
            if L and R:
                return node
            else:
                return L or R
            
        def dist(node, target):
            if not node:
                return sys.maxsize # initial min
            if node.val == target:
                return 0
            return min(dist(node.left, target), dist(node.right, target)) + 1
        
        lca = postorder(root) # same as LC 236, literally the classic LCA Problem
        return dist(lca, p) + dist(lca, q)