Algorithm for Finding the Diameter of a Tree¶

Using DP¶

Root the tree at vertex 1. Using dynamic programming, we compute the following two variables for each vertex \(u\):

• \(dp_u\) represents the maximum distance from vertex \(u\) to any vertex within its subtree.

• \(ans_u\) represents the diameter of the subtree rooted at \(u\).

Clearly, the final answer to the problem is \(ans_1\). The remaining task is to determine how to compute these two variables.

Computing \(dp_u\): Observe that the first move from \(u\) must be to one of its children. Therefore, we choose the child with the maximum \(dp\) value. Thus, \(dp_u = 1 + \max(\text{dp of children})\).

Computing \(ans_u\): We consider two cases based on whether vertex \(u\) lies on the diameter or not:

• If \(u\) is not on the diameter, then \(ans_u\) equals the maximum \(ans\) among its children (since the diameter must be entirely within one of the subtrees).

• If \(u\) is an endpoint of the diameter, the answer would be \(dp_u\).

• If \(u\) lies on the path between two endpoints, we consider two children \(a\) and \(b\) that contribute to this path. The diameter becomes \(2 + dp_a + dp_b\). Therefore, we select the two children with the largest \(dp\) values.

In the code below, \(mx1\) and \(mx2\) store the first and second largest \(dp\) values among the children, respectively.

const int maxn = 1e5 + 10;

vector<int> g[maxn];
int dp[maxn], ans[maxn];

void dfs(int u, int par = 0){
    int mx1 = -1, mx2 = -1;
    for(int y : g[u]){
        if(y != par){
            dfs(y, u);
            dp[u] = max(dp[u], 1 + dp[y]);
            ans[u] = max(ans[u], ans[y]);
            if(mx1 == -1 || dp[mx1] < dp[y]){
                mx2 = mx1;
                mx1 = y;
            }
            else if(mx2 == -1 || dp[mx2] < dp[y])
                mx2 = y;
        }
    }
    ans[u] = max(ans[u], dp[u]);
    if(mx1 != -1 && mx2 != -1){
    ans[u] = max(ans[u], 2 + dp[mx1] + dp[mx2]);
}

dfs up/down¶

Sometimes our goal is to compute a variable like \(dp\) for each node of a tree, but calculating \(dp_u\) requires knowing the \(dp\) values of all adjacent nodes of \(u\) (not just its children).

The simplest example to introduce this technique is the problem of finding the maximum distance from each node. Suppose we want to compute the eccentricity of each node \(u\), which we'll call \(ans_u\). To find the answer for a single node, we can simply root the tree at that node and compute its height in \(O(n)\). But can we solve the problem for all nodes together in \(O(n)\)?

The first challenge is that since calculating the answer for a node depends on its neighbors' answers, we don't know where to start the computation!

Root the tree at node \(u\). Breaking the problem into two parts can be helpful. Let \(dpDown_u\) be the maximum distance from node \(u\) to any node in its subtree. Also, let \(dpUp_u\) be the maximum distance from node \(u\) to any node outside its subtree (i.e., we must first move to \(u\)'s parent). Clearly, the answer for node \(u\) is the maximum of \(dpDown_u\) and \(dpUp_u\).

As discussed earlier, \(dpDown_u\) can be computed from the \(dpDown\) values of \(u\)'s children.

To compute \(dpUp_u\), note that after moving from \(u\) to its parent \(par\), we have two options:

1. Move further upward: In this case, the answer is \(1 + dpUp_{par}\) (assuming \(par\) is the parent of \(u\)).

2. Move downward: Go to one of \(u\)'s siblings, say \(w\), then descend. Here, the answer is \(2 + dpDown_w\).

The key insight is that we don't need to check all siblings of \(u\) every time to find the sibling \(w\) with the maximum \(dpDown\). For each parent \(par\), it's sufficient to precompute the two children with the largest \(dpDown\) values. The node \(w\) will always be one of these two children. (Why?)

Now we understand how to compute the variables, but an unresolved issue remains: What order should we follow for computation? Calculating \(dpDown\) requires values from children, while calculating \(dpUp\) requires values from the parent. Which should we compute first?

The answer is simple yet clever: We can compute the values in two phases. First, compute all \(dpDown\) values using a dfsDown, then compute all \(dpUp\) values using a dfsUp! The key is that in dfsDown, children's values are computed before the current node's value. In dfsUp, however, the parent's value is computed first, and then children's values are derived from the parent.

Note that in the dfsUp function, when processing a node, we assume its \(dpUp\) value has already been computed, and then we compute the \(dpUp\) values of its children.

const int maxn = 1e5 + 10;

vector<int> g[maxn];
int dpUp[maxn], dpDown[maxn];

void dfsDown(int u, int par = 0){ // aval bayad in taabe ra ejra konim
    for(int y : g[u]){
        if(y != par){
            dfsDown(y, u);
            dpDown[u] = max(dpDown[u], dpDown[y] + 1);
        }
    }
}
void dfsUp(int u, int par = 0){
   int mx1 = -1, mx2 = -1;
   for(int y : g[u]){
       if(y != par){
           if(mx1 == -1 || dpDown[mx1] < dpDown[y]){
                mx2 = mx1;
                mx1 = y;
           }
           else if(mx2 == -1 || dpDown[mx2] < dpDown[y]){
                mx2 = y;
           }
       }
   }
   for(int y : g[u]){
       if(y != par){
            if(y == mx1){
                dpUp[y] = dpUp[u] + 1;
                if(mx2 != -1)
                    dpUp[u] = max(dpUp[u], doDown[mx2] + 2);
            }
            else{
                dpUp[y] = max(dpUp[u]+1, doDown[mx1] + 2);
            }
            dfsUp(y, u);
       }
   }
}

A Simpler Algorithm¶

The diameter of a tree has a property that helps us find it more simply. This property is: the farthest vertex from any vertex is one end of a diameter of the tree.

To prove this, root the tree from an arbitrary vertex. Consider one of the diameters of the tree (which is also a path). This diameter/path has exactly one vertex that is closest to the root (the lowest common ancestor of the path's endpoints). If the farthest vertex from our initial vertex lies within the subtree of this common ancestor, we can remove one branch and add the branch corresponding to this vertex. The path length does not decrease, so it remains a diameter. The other case is impossible because starting from a vertex farther from the common ancestor and moving to this vertex would create a longer diameter, which is a contradiction. Thus, this vertex must be one end of a diameter of the tree.

Using this property, we can find the tree's diameter. We write a function that takes a vertex as input and returns one of the farthest vertices from it using BFS. We run this function from an arbitrary vertex and call the result \(u\). Then, we run the same function again starting from \(u\) and call the result \(v\). From the above theorem, since \(u\) is one end of a diameter, the path \(uv\) is a diameter of the tree.