Tree Diameter CSES

Problems and Contests
#1

Link to the problem

#include <bits/stdc++.h>

using namespace std;
using pi = pair<int, int>;

vector<int> adj[200000];
vector<pi> nodePos;
bool visited[200000];

int root;

void DFS(int curNode, int component, int depth){
    if(visited[curNode]){
        return;
    }
    visited[curNode] = true;
    nodePos.push_back({depth, component});
    for(int adjNode: adj[curNode]){
        if(curNode == root){
            DFS(adjNode, ++component, depth+1);
        } else {
            DFS(adjNode, component, depth+1);
        }
    }
}

int main(){
    int n; cin >> n;
    for(int i=0; i<n-1; i++){
        int a, b; cin >> a >> b;
        --a, --b;
        adj[a].push_back(b);
        adj[b].push_back(a);
    }
    int maxConnected = 0;
    root = 0;
    for(int i=0; i<n; i++){
        if(adj[i].size() > maxConnected){
            root = i;
            maxConnected = adj[i].size();
        }
    }
    DFS(root, 0, 0);
    sort(nodePos.begin(), nodePos.end(), [](pi a, pi b){
        if(a.first == b.first){
            return a.second > b.second;
        } else {
            return a.first > b.first;
        }
    });
    int ans = 0;
    ans += nodePos[0].first;
    for(int i=1; i<n; i++){
        if(nodePos[i].second != nodePos[0].second){
            ans += nodePos[i].first;
            break;
        }
    }
    cout << ans << endl;
}

My logic was to find the two largest tree depths with respect to a given root node and find the two maximum tree depths with different “components” in the tree. There is one edge case that is breaking my solution. Can someone help me identify what it is? TIA.

#2
#3

Do I need to rewrite my solution from start @Benq?

#4

Your implementation seems different from either of the two standard algorithms for this problem:

  • dynamic programming with finding two greatest depths per subtree
  • find farthest node from any node A = B, then find farthest node from B = C. Answer is distance from B to C.

The DP solution should only be a few lines of code in the DFS. See https://codeforces.com/blog/entry/20935 for DP on trees (first example is diameter).