Problem Description

You are given an undirected tree with n nodes labeled from 0 to n - 1. The tree is represented by a 2D array edges where edges[i] = [ai, bi] indicates there is an undirected edge between nodes ai and bi. The array has exactly n - 1 edges (which makes it a valid tree).

The diameter of a tree is defined as the number of edges in the longest path between any two nodes in the tree.

Your task is to find and return the diameter of the given tree.

For example, if you have a tree with edges connecting nodes in a linear fashion like 0-1-2-3, the longest path would be from node 0 to node 3, which contains 3 edges, so the diameter would be 3.

The key insight is that the longest path in a tree (the diameter) will always have its endpoints at two "peripheral" nodes - nodes that are as far apart as possible from each other in the tree structure.

Flowchart Walkthrough

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

Is it a graph?

• Yes: The problem explicitly states we have an undirected tree with nodes and edges. A tree is a special type of graph (connected and acyclic).

Is it a tree?

• Yes: The problem clearly mentions it's an undirected tree with n nodes and n-1 edges, which is the exact definition of a tree structure.

DFS

• According to the flowchart, when we have a tree structure, we should use DFS (Depth-First Search).

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

This makes perfect sense because:

1. Finding the diameter requires exploring paths in the tree, which DFS handles naturally by traversing as deep as possible along each branch
2. The solution approach of running DFS twice (first to find one endpoint of the diameter, then from that endpoint to find the other end) aligns perfectly with DFS's ability to explore tree paths exhaustively
3. DFS allows us to track the distance/depth as we traverse, which is essential for finding the longest path (diameter)

Intuition

The key insight for finding the tree diameter comes from understanding what the longest path in a tree looks like. The longest path must connect two "leaf-like" nodes that are as far apart as possible - these are the peripheral nodes of the tree.

Think about it this way: if we pick any arbitrary node and find the farthest node from it, that farthest node must be one of the endpoints of the diameter. Why? Because if it weren't, then there would be an even farther node, which contradicts our finding.

Here's the clever observation:

1. Start from any node (let's say node 0) and perform DFS to find the farthest node from it. Let's call this node a.
2. Node a is guaranteed to be one endpoint of the diameter path.
3. Now perform another DFS from node a to find the farthest node from a. This gives us node b.
4. The path from a to b is the diameter of the tree.

This works because in a tree:

• There's exactly one path between any two nodes
• The farthest node from any starting point will always be at one of the "extremes" of the tree
• The two nodes that are maximally distant from each other define the diameter

The algorithm essentially "stretches" the tree twice - first to find one end, then to find the other end from that first endpoint. This two-pass DFS approach guarantees we find the true diameter without having to check all possible pairs of nodes, which would be inefficient.

The beauty of this solution is its simplicity - we only need two DFS traversals instead of checking all O(n²) possible paths, making it an O(n) solution.

Learn more about Tree, Depth-First Search, Breadth-First Search, Graph and Topological Sort patterns.

Solution Approach

The implementation follows the two-DFS approach we identified:

Step 1: Build the Graph Representation We first convert the edge list into an adjacency list using a defaultdict(list). For each edge [a, b], we add b to a's neighbor list and a to b's neighbor list since the tree is undirected.

g = defaultdict(list)
for a, b in edges:
    g[a].append(b)
    g[b].append(a)

Step 2: DFS Implementation The DFS function takes three parameters:

• i: current node being visited
• fa: parent node (to avoid revisiting in an undirected tree)
• t: current distance from the starting node

def dfs(i: int, fa: int, t: int):
    for j in g[i]:
        if j != fa:  # Don't go back to parent
            dfs(j, i, t + 1)
    nonlocal ans, a
    if ans < t:  # Update if we found a farther node
        ans = t
        a = i

The DFS explores all neighbors except the parent node. It tracks the maximum distance (ans) and the corresponding farthest node (a).

Step 3: Two-Pass DFS Execution

1. First DFS: Start from node 0 with parent -1 and distance 0. This finds the farthest node from node 0, which becomes our first endpoint stored in variable a.

ans = a = 0
dfs(0, -1, 0)

2. Second DFS: Start from the node a found in the first pass. This finds the farthest node from a, and the distance to this node is the diameter of the tree.

dfs(a, -1, 0)
return ans

Why This Works:

• The first DFS guarantees that node a is one endpoint of the diameter
• The second DFS from a finds the other endpoint, and the distance between them is the maximum possible in the tree
• The fa parameter prevents cycles during traversal in the undirected tree
• The nonlocal keyword allows the nested DFS function to modify the outer scope variables ans and a

Time Complexity: O(n) - We visit each node at most twice (once in each DFS)
Space Complexity: O(n) - For the adjacency list and recursion stack

Example Walkthrough

Let's walk through a small example to illustrate the solution approach.

Consider a tree with 5 nodes and the following edges: [[0,1], [1,2], [1,3], [3,4]]

This creates the following tree structure:

        0
        |
        1
       / \
      2   3
          |
          4

Step 1: Build the adjacency list

g[0] = [1]
g[1] = [0, 2, 3]
g[2] = [1]
g[3] = [1, 4]
g[4] = [3]

Step 2: First DFS from node 0 Starting from node 0 with distance 0:

• Visit node 0 (distance=0)
  • Visit node 1 (distance=1)
    • Visit node 2 (distance=2) → Update: ans=2, a=2
    • Visit node 3 (distance=2) → No update (ans already 2)
      • Visit node 4 (distance=3) → Update: ans=3, a=4

After first DFS: ans=3, a=4 (node 4 is farthest from node 0)

Step 3: Second DFS from node 4 Reset ans to 0 and start from node 4:

• Visit node 4 (distance=0)
  • Visit node 3 (distance=1)
    • Visit node 1 (distance=2)
      • Visit node 0 (distance=3) → Update: ans=3, a=0
      • Visit node 2 (distance=3) → No update (ans already 3)

After second DFS: ans=3

The diameter is 3, which represents the longest path from node 4 to either node 0 or node 2 (both paths have 3 edges: 4→3→1→0 or 4→3→1→2).

This demonstrates how the two-DFS approach correctly identifies:

1. One endpoint of the diameter (node 4) in the first pass
2. The actual diameter length (3 edges) in the second pass from that endpoint

Solution Implementation

Time and Space Complexity

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

The algorithm performs two depth-first searches (DFS) on the tree:

• First DFS: Traverses all nodes once starting from node 0 to find the farthest node from it, visiting each node exactly once - O(n)
• Second DFS: Traverses all nodes once starting from the farthest node found in the first DFS, visiting each node exactly once - O(n)
• Building the adjacency list from edges takes O(n-1) = O(n) time since a tree with n nodes has n-1 edges

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

Space Complexity: O(n)

The space usage consists of:

• Adjacency list g: Stores all edges in both directions, requiring O(2*(n-1)) = O(n) space
• Recursion call stack: In the worst case (a skewed tree), the DFS recursion depth can be O(n)
• Variables ans and a: O(1) constant space

Total space complexity: O(n) dominated by the adjacency list and recursion stack.

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

Common Pitfalls

1. Handling Empty Edge List

The code assumes there's at least one node in the tree. If edges is empty (single node tree with no edges), the code will still work but it's worth noting that the diameter would be 0.

Pitfall Example:

edges = []  # Single node tree
# The graph will be empty, but dfs(0, -1, 0) might fail if node 0 doesn't exist

Solution: Add an edge case check:

if not edges:
    return 0

2. Incorrect Parent Tracking in DFS

A common mistake is forgetting to pass the parent node correctly or using the wrong initial parent value. Using 0 instead of -1 as the initial parent could cause issues if node 0 has neighbors.

Pitfall Example:

# Wrong: Using 0 as initial parent
dfs(0, 0, 0)  # This would skip node 0's actual neighbors if 0 is in their adjacency list

# Correct: Using -1 as initial parent
dfs(0, -1, 0)

3. Variable Scope Issues with Nested Functions

Forgetting to use nonlocal when modifying outer scope variables in the nested DFS function will create local variables instead, causing the algorithm to fail.

Pitfall Example:

def dfs(node, parent, distance):
    # Missing nonlocal declaration
    if max_distance < distance:  # This creates a local variable
        max_distance = distance   # Won't update the outer scope variable

Solution: Always declare outer scope variables with nonlocal:

def dfs(node, parent, distance):
    nonlocal max_distance, farthest_node
    if max_distance < distance:
        max_distance = distance

4. Not Resetting Distance Between DFS Calls

Forgetting to reset max_distance to 0 before the second DFS will cause incorrect results since the variable still holds the value from the first DFS.

Pitfall Example:

# First DFS
dfs(0, -1, 0)
# Forgot to reset max_distance here!
# max_distance still holds the value from first DFS
dfs(farthest_node, -1, 0)  # Won't update correctly

Solution: Always reset the distance before the second DFS:

dfs(0, -1, 0)
max_distance = 0  # Reset before second DFS
dfs(farthest_node, -1, 0)

5. Starting DFS from a Non-Existent Node

If the tree nodes aren't labeled from 0 to n-1 consecutively, starting DFS from node 0 might fail.

Pitfall Example:

edges = [[1, 2], [2, 3]]  # No node 0 exists
# dfs(0, -1, 0) will traverse nothing since node 0 has no neighbors

Solution: Start from any existing node in the graph:

if edges:
    start_node = edges[0][0]  # Use any node from edges
    dfs(start_node, -1, 0)