In this blog, I will go over a clean way to find the diameter of a tree (and recover the nodes on the diameter). If you are unsure of what the diameter of a tree means, consult this blog.

The standard algorithm for finding the endpoints of a diameter is as follows:

1. Pick an arbitrary node (such as node $$$0$$$) and find the farthest node from it using a DFS. Let's call this new node $$$a$$$.

2. Find the farthest node from $$$a$$$ with another DFS. Let's call this new node $$$b$$$.

The path from $$$a$$$ to $$$b$$$ is a diameter of a tree. Note that there are multiple valid diameters in a tree.

Implementations of this algorithm use a DFS that returns the deepest leaf in a node's subtree and the depth of that leaf. A typical implementation in C++ is below.

auto dfs = [&](int u, int p, auto &&dfs) -> pair<int, int> {
    pair<int, int> res{-1, u};
    for(int v : adj[u]) if(v != p)
        res = max(res, dfs(v, u, dfs));
    res.first++;
    return res;
};
auto [dist1, a] = dfs(0, -1, dfs);
auto [dist2, b] = dfs(a, -1, dfs);

Here, dist2 is also the length of the tree diameter (in terms of edges).

Sometimes, you want the nodes on the diameter path. You recover the path by storing the parents of nodes and recovering the path with a separate loop:

vector<int> parent(n, -1);
auto dfs = [&](int u, int p, auto &&dfs) -> pair<int, int> {
    parent[u] = p;
    pair<int, int> res{-1, u};
    for(int v : adj[u]) if(v != p)
        res = max(res, dfs(v, u, dfs));
    res.first++;
    return res;
};
auto [dist1, a] = dfs(0, -1, dfs);
auto [dist2, b] = dfs(a, -1, dfs);
vector<int> diameter;
for(int x = b; x != -1; x = parent[x])
    diameter.push_back(x);

However, there is a funnier (and a bit shorter) way of accomplishing the same thing. What if we return the path to the deepest leaf instead of returning the deepest leaf and its depth?

auto dfs = [&](int u, int p, auto &&dfs) -> vector<int> {
    vector<int> best;
    for (int v : adj[u]) if (v != p) {
        auto cur = dfs(v, u, dfs);
        if (size(cur) > size(best)) swap(cur, best);
    }
    best.push_back(u);
    return best;
};
auto diameter = dfs(0, -1, dfs);
diameter = dfs(diameter[0], -1, dfs);

The code finds the path to the deepest leaf in a subtree and adds the current node to the path. The time complexity appears to be $$$O(n^2)$$$ on trees with large diameters (such as a line), but it is actually $$$O(n)$$$. Due to move semantics and return value optimization, the returned vector will be moved in $$$O(1)$$$ instead of copied in $$$O(size)$$$.

If you would like to try out these implementations or your own, use this problem. Please share your own implementations if you have alternate ideas.