PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: progokcoe
Tester: mexomerf
Editorialist: iceknight1093
DIFFICULTY:
TBD
PREREQUISITES:
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
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.