EXEPDIAM - Editorial

PROBLEM LINK:

Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4

Author: progokcoe
Tester: mexomerf
Editorialist: iceknight1093

DIFFICULTY:

TBD

PREREQUISITES:

Diameter of a tree

PROBLEM:

You have a tree with N vertices. You want to attach two more vertices to it.
Find the number of ways to do this such that the diameter of the new tree is strictly larger than the original diameter.

EXPLANATION:

Let D denote the diameter of the original tree (which can be found using DFS, see the link in prerequisites section if you don’t know how).

Let the two new nodes be u_1 and u_2.
There are three possibilities for attaching them to the tree:

  • u_1 and u_2 can be attached to two different vertices; or
  • u_1 and u_2 can be attached to the same vertex; or
  • u_1 can be attached to some vertex and then u_2 attached to it, or vice versa.

We count the possibilities for all three cases separately.

Case 1: Different vertices

Suppose u_1 is attached to v_1, and u_2 is attached to v_2.

Then, the only way for the diameter to increase is if either v_1 or v_2 are themselves endpoints of some diameter of the original tree.

So, suppose there are X vertices in total that are endpoints of some diameter.
Then, the number of ways to choose v_1 and v_2 is

X\cdot (X-1) + 2X\cdot (N-X)

because either both v_1 and v_2 are (different) diameter endpoints, or one of them is and the other isn’t.

This requires us to find the value of X, but if you’re familiar with diameters that isn’t very hard.
Let \text{len}[v] denote the length of the longest path with v as one of its endpoints.
Then, if (x, y) is a diameter of the tree, \text{len}[v] = \max(d(v, x), d(v, y)).
v is an endpoint of some diameter if and only if \text{len}[v] = D.

The diameter endpoints x and y were obtained for free when we first found the diameter, and computing distance from them to everything else is easy with DFS so X can be found in linear time.

Case 2: Same vertex

Suppose u_1 and u_2 are attached to v.
Once again, it can be proved that this results in the diameter increasing if and only if v is a diameter endpoint in the original tree; and we already know the count of those from the first case.

Case 3: Chained together

Suppose u_1 is attached to v, and u_2 is attached to u_1.
This increases the diameter if and only if:

  • v is a diameter endpoint; or
  • \text{len}[v] = D-1 (which would make the diameter D+1).

We know the \text{len}[v] values so just count the valid ones and add them to the answer.
Remember to multiply by two, to account for the case where u_2 is attached first.

TIME COMPLEXITY:

\mathcal{O}(N) per testcase.

CODE:

I did the same approach then why is my code wrong

vector<int> adj[maxN];
vector<int>dist;

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;
}

Can u explain why did u do dfs with the v again?

Why is the code to the problem missing in the editorial?

My solution: https://www.codechef.com/viewsolution/1041119392
I seem to be having trouble with finding the count of number of end vertices of diameter and/or the number of vertices where maximum length to other vertex is one less than diameter.
Any help with rectifying the count obtained is appreciated.

1 Like