Tree diameter: two-BFS to find the longest path

Coding Challenges
, , ,
1

Problem. Given an unrooted tree with n nodes and n−1 edges, find the diameter — the number of edges on the longest path between any two nodes.

Examples:

  • n=5, edges=[(0,1),(1,2),(2,3),(1,4)] → 3 — path 0→1→2→3.
  • n=1, edges=[] → 0 — single node.
  • n=2, edges=[(0,1)] → 1
  • n=7, edges=[(0,1),(1,2),(1,3),(3,4),(3,5),(5,6)] → 4 — path 2→1→3→5→6.

Solution. Two-BFS: BFS from any node to find the farthest node u; BFS from u to find the farthest node v. The distance u→v is the diameter.

from collections import deque


def tree_diameter(n, edges):
    if n <= 1:
        return 0
    adj = [[] for _ in range(n)]
    for u, v in edges:
        adj[u].append(v)
        adj[v].append(u)

    def bfs_farthest(start):
        dist = [-1] * n
        dist[start] = 0
        q = deque([start])
        farthest = start
        while q:
            node = q.popleft()
            for nb in adj[node]:
                if dist[nb] == -1:
                    dist[nb] = dist[node] + 1
                    q.append(nb)
                    if dist[nb] > dist[farthest]:
                        farthest = nb
        return farthest, dist[farthest]

    u, _ = bfs_farthest(0)
    _, diameter = bfs_farthest(u)
    return diameter

O(n) time, O(n) space.

Gotchas.

  • Works only on trees (connected, acyclic). For general graphs, use all-pairs shortest paths.
  • A single node has diameter 0, not 1.
  • For weighted edges, replace BFS with Dijkstra (or BFS on unweighted after edge relaxation).
  • The two-BFS trick relies on the tree property: the farthest node from any node is always an endpoint of a diameter.

Run: tree diameter