Solutions

# At each node compute the left and right heights, then the diameter as their sum. Use a global variable to update the maximum.
# trees

# Recursive DFS - O(n) time, O(h) space
def diameter_of_binary_tree(root: TreeNode) -> int:
    res = 0

    def dfs(root):
        nonlocal res
        
        if not root:
            return 0
        
        left_height = dfs(root.left)
        right_height = dfs(root.right)

        res = max(left_height + right_height, res)
        return 1 + max(left_height, right_height)
    
    dfs(root)
    return res


# Iterative DFS - O(n) time, O(n) space
def diameter_of_binary_tree(root: TreeNode) -> int:
    stack = [root]
    mp = {None: (0, 0)}

    while stack:
        node = stack[-1]

        if node.left and node.left not in mp:
            stack.append(node.left)
        elif node.right and node.right not in mp:
            stack.append(node.right)
        else:
            node = stack.pop()

            leftHeight, leftDiameter = mp[node.left]
            rightHeight, rightDiameter = mp[node.right]

            mp[node] = (1 + max(leftHeight, rightHeight),
                        max(leftHeight + rightHeight, leftDiameter, rightDiameter))

    return mp[root][1]