https://www.codechef.com/problems/EXEPDIAM
Code
vector adj[maxN];
vectordist;
void dfs(int u , int p)
{
for (auto v : adj[u])
{
if (v == p) continue;
dist[v] = dist[u] + 1;
dfs(v , u);
}
}
void solve()
{
int n ; cin >> n;
for (int i = 0 ; i < n ; i++) adj[i].clear();
for (int i = 0 ; i < n - 1 ; i++)
{
int u , v ; cin >> u >> v;
u--; v--;
adj[u].push_back(v);
adj[v].push_back(u);
}
if (n == 1)
{
cout << 3 << endl;
return;
}
dist.assign(n , 0);
dfs(0 , -1);
int u = max_element(all(dist)) - dist.begin();
dist.assign(n , 0);
dfs(u , -1);
int v = max_element(all(dist)) - dist.begin();
int max_d = *max_element(all(dist));
set<int>st1 , st2;
for (int i = 0 ; i < n ; i++)
{
if (i == u)continue;
if (dist[i] == max_d)st1.insert(i);
if (dist[i] == max_d - 1)st2.insert(i);
}
dist.assign(n , 0);
dfs(v , -1);
for (int i = 0 ; i < n ; i++)
{
if (i == v)continue;
if (dist[i] == max_d)st1.insert(i);
if (dist[i] == max_d - 1)st2.insert(i);
}
int ans = 2 * st2.size();
ans += (2 * (max_d) * st1.size());
int x = st1.size();
ans += (x * (x - 1));
ans += x;
cout << ans << endl;
}








Auto comment: topic has been updated by ONE___PIECE (previous revision, new revision, compare).
Auto comment: topic has been updated by ONE___PIECE (previous revision, new revision, compare).
First find the diameter of the tree (either dfs or DP) let's call it dist, then count the number of nodes such that the maximum distance from it to any other node is dist or dist — 1 (you can do it by rotating root DP), let's call it cnt1 and cnt2. Then the answer is just cnt1 * (2 + 2 * n — cnt1) + 2 * cnt2.