See all coding puzzles and patterns here.

Problem Definition

Given a binary tree, determine its diameter, the maximum length of the edges between any two nodes in the tree.


Node Class

The class representing a node in the tree:

class Node:
    def __init__(self, value, left=None, right=None):
        self.value = value
        self.left = left
        self.right = right

Examples

For this tree:
    x
   / \
  x   x
 / \
x   x
 \   \
  x   x
 /     \
a       x
         \
          b
The answer is 7, the number of edges between a and b.

Solution

We'll use a depth-first traversal to find the maximum length up to each sub-node as well as the maximum length inside the node's tree. We can then use these values within the stack frame processing a parent node to compute the maximum length up to the parent node as well as the maximum length inside the node's tree.

def helper(root: Node) -> tuple[int, int]:
    """
    Returns (max length of tree, max length up to root)
    """
    ans = 0
    if root.left is None:
        left = 0
    else:
        ans_left, left = helper(root.left)
        left += 1
        ans = ans_left
    if root.right is None:
        right = 0
    else:
        ans_right, right = helper(root.right)
        right += 1
        ans = max(ans_left, ans)
    return max(ans, left + right), max(left, right)

def diameter(root: Node) -> int:
    return helper(root)[0]

Complexity

Let $n = $ the number of nodes in the tree.

Since we're traversing each node exactly once, the runtime complexity should be $\Theta(n)$.

Since we're generating a stack frame for each node and since in the worst case the tree might be a linked list, the auxiliary space complexity should be $O(n)$.