Lowest Common Ancestor of a Binary Tree

Patrick Leaton
Problem Description

Given a binary tree, find the lowest common ancestor (LCA) of two given nodes in the tree.

According to the definition of LCA on Wikipedia: “The lowest common ancestor is defined between two nodes p and q as the lowest node in T that has both p and q as descendants (where we allow a node to be a descendant of itself).”

 

Example 1:

Input: root = [3,5,1,6,2,0,8,null,null,7,4], p = 5, q = 1
Output: 3
Explanation: The LCA of nodes 5 and 1 is 3.

Example 2:

Input: root = [3,5,1,6,2,0,8,null,null,7,4], p = 5, q = 4
Output: 5
Explanation: The LCA of nodes 5 and 4 is 5, since a node can be a descendant of itself according to the LCA definition.

Example 3:

Input: root = [1,2], p = 1, q = 2
Output: 1

 

Constraints:

  • The number of nodes in the tree is in the range [2, 10^5].
  • -10^9 <= Node.val <= 10^9
  • All Node.val are unique.
  • p != q
  • p and q will exist in the tree.

 

The description was taken from https://leetcode.com/problems/lowest-common-ancestor-of-a-binary-tree/.

Problem Solution

#O(N) Time, O(N) Space
class Solution:
    def lowestCommonAncestor(self, root: 'TreeNode', p: 'TreeNode', q: 'TreeNode') -> 'TreeNode':
        output = None
 
        def dfs(node:TreeNode) -> int:
            if not node:
                return 0
            nonlocal output
            left = dfs(node.left)
            right = dfs(node.right)
            itself = node == p or node == q
            if itself + left + right == 2:
                output = node
                return 1
            return left or right or itself
 
        dfs(root)
        return output

Problem Explanation


In this question, we will have two scenarios for the lowest common ancestor.

Either we will have two nodes that come from two different branches of the family tree, in which case the node that splits the two branches that house p and q would be the lowest common ancestor, like in example one.

Or, we will have p and q within the same branch of the family tree.  Then, whichever of the two nodes further up the family tree would be the lowest common ancestor between the two, like in example two.

Similar to the precursor to this question, Lowest Common Ancestor of a Binary Search Tree, we will perform a Depth-First Search.  During the Depth-First Search within each DFS call, we will assign a value, zero or one.  Zero will be for null nodes and nodes that have neither one of the input nodes attached to one of their child branches.  One will be for a node that is either p or q

If a DFS returns two, then that means that we have found a common ancestor between the two input nodes, and the first common ancestor found will be the lowest common ancestor by default due to the nature of Depth-First Search.


Let's start by initializing our output node to none.

        output = None


Then, we will make our recursive dfs function.

For binary tree questions, one key question to ask is "whose information do we need first?".  For anything that requires heights or things to that degree, we will need a postorder traversal because, in order for a node to know its height, it would have to know the heights of each child node beneath it.  In this question, we are looking for the lowest common ancestor, so that means we will have to visit each node after its children to know whether a node has a child node that it is a common ancestor of.

        def dfs(node:TreeNode) -> bool:

 

Our base case will be if we have reached a node with the value of none, then we have exhausted a DFS path so we will return zero since this node is neither p nor q.

            if not node:
                return 0

 

Then, we will need to declare that the output variable is not local to this dfs function, as it belongs to the parent function.

            nonlocal output

 

Afterward, we will recurse the function by passing in the left and right children.

            left = dfs(node.left)
            right = dfs(node.right)

 

A DFS goes down until it hits a base case, in this case, a null node, then it returns back up like a yoyo.

At this point of the function, a DFS has returned from the left and right children that was recursed in the previous two lines in which case, we will check the current node as being either p or q.

            itself = node == p or node == q

 

If the current node, the left path, and the right path combined have a total sum of two, then we have found both p and q in the current DFS path and this is the first time that we have seen both nodes required for the current node to be a common ancestor.  In this case, we will declare the output as the current node and return one, since the above calculation will be done in parent calls and will need a number value returned to avoid errors.

            if itself + left + right == 2:
                output = node
                return 1

 

If the sum was one, then we have only seen one node and not the other.  If the sum is three, then that means we have declared the lowest common ancestor within a deeper DFS path.

If we don't have two, we will return that either the left path, the current node, or the right path were marked with having p or q so that a parent function call can perform the previous calculation.

            return left or right or itself


Now that our dfs function is built, the only thing we need to do is make the initial DFS call on the root and return the lowest common ancestor in the output.

        dfs(root)
        return output