Problem Description

You are given a binary tree and two integer values p and q. Your task is to find the distance between the two nodes that contain these values.

The distance between two nodes in a tree is defined as the number of edges you need to traverse to get from one node to the other along the path connecting them.

For example, if you have a binary tree and need to find the distance between nodes with values p and q, you would:

1. Find the node containing value p
2. Find the node containing value q
3. Count the number of edges in the path connecting these two nodes

The solution approach involves two key steps:

1. Finding the Lowest Common Ancestor (LCA): The lca function identifies the lowest common ancestor of the two target nodes. This is the deepest node in the tree that has both p and q as descendants (or is itself one of the target nodes).

2. Calculating distances from LCA: The dfs function computes the distance from a given node to a target value. Once we have the LCA, we calculate:

   • The distance from the LCA to node p
   • The distance from the LCA to node q
   • The total distance is the sum of these two distances

The algorithm works because the shortest path between any two nodes in a tree must pass through their lowest common ancestor. By finding this ancestor and measuring the distances from it to each target node, we get the total distance between the two nodes.

Flowchart Walkthrough

First, let's pin down the algorithm using the Flowchart. Here's a step-by-step walkthrough:

Is it a graph?

• Yes: A binary tree is a special type of graph where each node has at most two children and there are no cycles.

Is it a tree?

• Yes: The problem explicitly states we're working with a binary tree structure.

DFS

• Yes: We arrive at DFS (Depth-First Search) as the recommended approach.

Conclusion: The flowchart suggests using DFS for this binary tree problem.

This makes perfect sense for finding the distance between two nodes in a binary tree because:

1. We need to traverse the tree to find the lowest common ancestor (LCA)
2. We need to search from the LCA down to each target node to calculate distances
3. DFS allows us to efficiently explore paths in the tree recursively

The solution implements DFS in two ways:

• The lca function uses DFS to recursively find the lowest common ancestor by exploring left and right subtrees
• The dfs function uses DFS to find the distance from a given node to a target value by recursively exploring the tree and counting edges

Both operations naturally fit the DFS pattern of exploring as far as possible along each branch before backtracking, which is ideal for tree traversal problems.

Intuition

To find the distance between two nodes in a tree, we need to think about what the path between them looks like. In a tree, there's exactly one path between any two nodes (since trees have no cycles). This path must go up from one node to some common ancestor, then down to the other node.

The key insight is that every path between two nodes must pass through their Lowest Common Ancestor (LCA). The LCA is the deepest node that has both target nodes as descendants. Once we identify this pivot point, the problem becomes much simpler.

Think of it like finding the distance between two cities connected by roads that form a tree structure. If you want to travel from city p to city q, you might need to go up to a junction point (the LCA) and then down to your destination. The total distance is:

• Distance from p up to the LCA
• Plus distance from the LCA down to q

This naturally leads us to a two-phase approach:

1. First, find the LCA of nodes p and q using DFS. We recursively search the tree, and when we find both targets in different subtrees of a node, that node is our LCA.
2. Then, calculate the distance from the LCA to each target node separately using DFS again.

The beauty of this approach is that we're breaking down a complex path-finding problem into simpler subproblems: finding a common ancestor and measuring distances from that ancestor. Since the path from p to q = path from p to LCA + path from LCA to q, we can simply add these two distances together to get our answer.

Learn more about Tree, Depth-First Search, Breadth-First Search and Binary Tree patterns.

Solution Approach

The implementation consists of two main helper functions that work together to solve the problem:

1. Finding the Lowest Common Ancestor (lca function)

The lca function uses a recursive DFS approach to find the lowest common ancestor:

def lca(root, p, q):
    if root is None or root.val in [p, q]:
        return root
    left = lca(root.left, p, q)
    right = lca(root.right, p, q)
    if left is None:
        return right
    if right is None:
        return left
    return root

The algorithm works as follows:

• Base case: If we reach a None node or find a node with value p or q, we return that node
• Recursive search: We recursively search both left and right subtrees
• Decision logic:
  • If both left and right return non-null values, the current node is the LCA (both targets are in different subtrees)
  • If only one side returns a non-null value, that side contains both targets, so we propagate that result up
  • This ensures we find the deepest common ancestor

2. Calculating Distance from a Node (dfs function)

The dfs function calculates the distance from a given node to a target value:

def dfs(root, v):
    if root is None:
        return -1
    if root.val == v:
        return 0
    left, right = dfs(root.left, v), dfs(root.right, v)
    if left == right == -1:
        return -1
    return 1 + max(left, right)

The algorithm works as follows:

• Base cases:
  • If the node is None, return -1 (target not found)
  • If we find the target value, return 0 (distance to itself)
• Recursive search: Search both left and right subtrees
• Distance calculation:
  • If both subtrees return -1, the target isn't in this subtree
  • Otherwise, add 1 to the maximum of left and right (only one will be non-negative since each value appears once)

3. Combining the Results

The main function ties everything together:

g = lca(root, p, q)
return dfs(g, p) + dfs(g, q)

• First, find the LCA of nodes with values p and q
• Then, calculate the distance from the LCA to p
• Calculate the distance from the LCA to q
• Sum these distances to get the total path length

The time complexity is O(n) where n is the number of nodes, as we potentially visit each node during the LCA search and distance calculations. The space complexity is O(h) where h is the height of the tree, due to the recursive call stack.

Example Walkthrough

Let's walk through finding the distance between nodes with values 5 and 6 in this binary tree:

        3
       / \
      5   1
     / \   \
    6   2   4

Step 1: Find the Lowest Common Ancestor (LCA)

Starting from root (3), we search for nodes with values 5 and 6:

• At node 3: Check left subtree (contains 5) and right subtree (contains neither)
  • Left subtree search at node 5: This node matches value 5, so return node 5
  • Right subtree search at node 1: Neither 5 nor 6 found in this subtree, returns None
• Since only left subtree returns a non-null value, we continue searching there

Actually, let's trace this more carefully:

• At node 3:
  • Search left subtree starting at node 5
    • Node 5 matches p=5, return node 5
  • Search right subtree starting at node 1
    • Node 1 doesn't match, search its children
    • No left child, right child is 4 (doesn't match)
    • Returns None
  • Left returns node 5, right returns None → propagate node 5 up

Wait, we need to search for both 5 AND 6. Let me redo this:

• At node 3:
  • Search left subtree for both 5 and 6
    • At node 5: This matches p=5, return node 5
    • But wait, we also need to check for q=6 in its subtrees
    • Left child of 5 is node 6, which matches q=6
  • Since node 5 has value 5 and its child has value 6, node 5 is the LCA

Actually, the LCA algorithm works differently. Let me trace it correctly:

• lca(3, 5, 6):
  • Node 3 is not 5 or 6
  • left = lca(5, 5, 6) → returns node 5 (matches p=5)
  • right = lca(1, 5, 6) → returns None (neither value found)
  • Since only left is non-null, return left (node 5)
• The LCA is node 5

Step 2: Calculate distance from LCA to p (value 5)

• dfs(5, 5):
  • Node 5 has value 5
  • Return 0 (distance to itself)

Step 3: Calculate distance from LCA to q (value 6)

• dfs(5, 6):
  • Node 5 has value 5, not 6
  • left = dfs(6, 6) → Node 6 matches, return 0
  • right = dfs(2, 6) → Node 2 doesn't match, no children, return -1
  • Since left = 0 and right = -1, return 1 + 0 = 1

Step 4: Sum the distances

Total distance = 0 + 1 = 1 edge

This makes sense! To go from node 5 to node 6, we traverse exactly one edge (5 → 6).

Solution Implementation

Time and Space Complexity

Time Complexity: O(n) where n is the number of nodes in the binary tree.

The algorithm consists of two main operations:

1. Finding the Lowest Common Ancestor (LCA): The lca function traverses the tree once in the worst case, visiting each node at most once. This takes O(n) time.
2. Finding distances from LCA to both nodes: The dfs function is called twice - once for node p and once for node q. Each call traverses a subtree rooted at the LCA, and in the worst case (when the LCA is the root and one target node is a leaf), it may visit up to O(n) nodes. However, since we're only traversing from the LCA downward to find each target node, and each node is visited at most once across both DFS calls, the total time is O(n).

Overall time complexity: O(n) + O(n) = O(n)

Space Complexity: O(h) where h is the height of the binary tree.

The space complexity is determined by the recursion call stack:

1. The lca function uses recursion that goes as deep as the height of the tree, requiring O(h) space.
2. The dfs function also uses recursion with maximum depth equal to the distance from the LCA to the farthest target node, which is at most O(h).

In the worst case (skewed tree), h = O(n), making the space complexity O(n). In the best case (balanced tree), h = O(log n), making the space complexity O(log n).

Learn more about how to find time and space complexity quickly.

Common Pitfalls

1. Assuming Node Values are Unique

The most critical pitfall in this solution is the assumption that node values are unique in the tree. The LCA function uses node.val in [p_val, q_val] to identify target nodes, which breaks down when duplicate values exist.

Problem Example:

      1
     / \
    2   3
   /     \
  4       2  <- duplicate value

If searching for nodes with values p=2 and q=4, the algorithm might incorrectly identify the first occurrence of 2 as one of the targets, leading to wrong distance calculations.

Solution: Instead of searching by values, pass the actual TreeNode references:

def findDistance(self, root: Optional[TreeNode], p: TreeNode, q: TreeNode) -> int:
    def find_lca(node, p_node, q_node):
        if node is None or node == p_node or node == q_node:
            return node
        # ... rest remains the same
  
    def find_depth(node, target_node):
        if node is None:
            return -1
        if node == target_node:  # Compare references, not values
            return 0
        # ... rest remains the same

2. Not Handling Edge Cases

The code doesn't explicitly handle several edge cases that could cause issues:

• Same node for both p and q: When p equals q, the distance should be 0
• Non-existent values: If p or q don't exist in the tree, the function may behave unexpectedly
• Single node tree: Though the algorithm handles this correctly, it's worth explicit consideration

Solution: Add validation checks:

def findDistance(self, root: Optional[TreeNode], p: int, q: int) -> int:
    # Handle same value case
    if p == q:
        return 0 if self.node_exists(root, p) else -1
  
    # Verify both values exist
    if not self.node_exists(root, p) or not self.node_exists(root, q):
        return -1
  
    # Proceed with normal algorithm
    lca_node = find_lca(root, p, q)
    return find_depth(lca_node, p) + find_depth(lca_node, q)

def node_exists(self, root, val):
    if not root:
        return False
    if root.val == val:
        return True
    return self.node_exists(root.left, val) or self.node_exists(root.right, val)

3. Inefficient Multiple Tree Traversals

The current approach traverses the tree multiple times:

• Once to find the LCA
• Once to find distance from LCA to p
• Once to find distance from LCA to q

Solution: Combine operations in a single traversal:

def findDistance(self, root: Optional[TreeNode], p: int, q: int) -> int:
    def helper(node):
        """
        Returns tuple: (found_p, found_q, distance)
        distance is -1 if not yet found both
        """
        if not node:
            return False, False, -1
      
        left_p, left_q, left_dist = helper(node.left)
        right_p, right_q, right_dist = helper(node.right)
      
        # Check if current node is p or q
        is_p = node.val == p
        is_q = node.val == q
      
        # Aggregate found status
        found_p = is_p or left_p or right_p
        found_q = is_q or left_q or right_q
      
        # If we already found the distance in a subtree, propagate it
        if left_dist != -1:
            return found_p, found_q, left_dist
        if right_dist != -1:
            return found_p, found_q, right_dist
      
        # If this is the LCA (both found for first time)
        if found_p and found_q:
            # Calculate distance based on where p and q are
            dist = 0
            if left_p and left_q:
                dist = self.get_distance_in_subtree(node.left, p, q)
            elif right_p and right_q:
                dist = self.get_distance_in_subtree(node.right, p, q)
            else:
                # They're in different subtrees or one is current node
                dist = self.depth_to_value(node.left, p if left_p else q) + \
                       self.depth_to_value(node.right, p if right_p else q) + \
                       (0 if is_p or is_q else 2)
            return found_p, found_q, dist
      
        return found_p, found_q, -1
  
    _, _, distance = helper(root)
    return distance

These pitfalls highlight the importance of clearly understanding problem constraints, handling edge cases, and considering optimization opportunities when implementing tree algorithms.