Range Diameter Sum
Contest: Codeforces Round 691 (Div. 1)
Difficulty rating: 3500
Tags: data structures, trees

You are given a tree with $$$n$$$ vertices numbered $$$1, \ldots, n$$$. A tree is a connected simple graph without cycles.

Let $$$\mathrm{dist}(u, v)$$$ be the number of edges in the unique simple path connecting vertices $$$u$$$ and $$$v$$$.

Let $$$\mathrm{diam}(l, r) = \max \mathrm{dist}(u, v)$$$ over all pairs $$$u, v$$$ such that $$$l \leq u, v \leq r$$$.

Compute $$$\sum_{1 \leq l \leq r \leq n} \mathrm{diam}(l, r)$$$.

Input format:
The first line contains a single integer $$$n$$$ ($$$1 \leq n \leq 10^5$$$) — the number of vertices in the tree.

The next $$$n - 1$$$ lines describe the tree edges. Each of these lines contains two integers $$$u, v$$$ ($$$1 \leq u, v \leq n$$$) — endpoint indices of the respective tree edge. It is guaranteed that the edge list indeed describes a tree.

Output format:
Print a single integer — $$$\sum_{1 \leq l \leq r \leq n} \mathrm{diam}(l, r)$$$.

Constraints:
Time limit: 7.0 s
Memory limit: 512.0 MB

Example input:
4
1 2
2 4
3 2

Example output:
10

Example input:
10
1 8
2 9
5 6
4 8
4 2
7 9
3 6
10 4
3 9

Example output:
224

Editorial:
Please bear with my formal style for this one as my hand-waving skills are too weak to explain this properly.
We'll use the "divide-and-conquer" strategy. Let us implement a recursive procedure $$$solve(l, r)$$$ that will compute the sum $$$\sum_{l \leq i \leq j \leq r} diam(i, j)$$$. It will do so as follows:
• If $$$l = r$$$, then answer is trivially $$$0$$$.
• Otherwise, choose $$$m = \lfloor(l + r) / 2 \rfloor$$$, and compute the following sum: $$$\sum_{i = l}^m \sum_{j = m + 1}^r diam(i, j)$$$.
• The only entries $$$diam(i, j)$$$ that are not accounted for satisfy either $$$l \leq i \leq j \leq m$$$ or $$$m + 1 \leq i \leq j \leq r$$$. Compute their total as $$$solve(l, m) + solve(m + 1, r)$$$, and return the total answer.
Of course, $$$\sum_{i = l}^m \sum_{j = m + 1}^r diam(i, j)$$$ is the tricky part. We'll need some preparation to compute it. To make arguments cleaner, let us introduce an extra vertex in the middle of every edge, subdividing it into two. We'll refer to new vertices as extra vertices, and the old vertices as proper vertices. Since all distances are doubled, we will have to divide the answer by two in the end.
We will define a circle $$$C(v, r)$$$ of radius $$$r$$$ centered at a vertex $$$v$$$ in the tree as the set of all vertices at distance at most $$$r$$$ from the vertex $$$v$$$. We also assume that if $$$r > 0$$$ there are at least two vertices in $$$C(v, r)$$$ exactly at distance $$$r$$$ from $$$v$$$, otherwise decrease $$$r$$$ accordingly.
For a set of vertices $$$S$$$ we define $$$Cover(S)$$$ as a circle of smallest radius containing all vertices of $$$S$$$. If we assume that all vertices of $$$S$$$ are proper, then $$$Cover(S)$$$ is unique. Indeed, after subdividing all distances among $$$S$$$ are even, thus in $$$Cover(S) = (v, r)$$$ we must have $$$r = diam(S) / 2$$$, and $$$v$$$ is the unique diameter midpoint. We can see that if $$$r > 0$$$, then $$$S$$$ has at least two vertices at distance $$$r$$$ from the covering circle center $$$v$$$, namely, the diameter endpoints.
Warning: math ahead. It you're not interested in proofs, only read the premise of Lemma 2.
Lemma 1. If a circle $$$C(v, r)$$$ contains two proper vertices $$$a, b$$$, then it also contains the midpoint $$$c$$$ on the path between $$$a, b$$$, and further $$$dist(v, c) \leq r - dist(a, b) / 2$$$.
Proof. Say, for the vertex $$$a$$$ we have $$$dist(v, a) \leq dist(v, c) + dist(c, a) = dist(v, c) + dist(a, b) / 2$$$. This inequality is strict: $$$dist(v, a) < dist(v, c) + dist(c, a)$$$ only when $$$v$$$ and $$$a$$$ are in the same subtree with respect to $$$c$$$. If we assume that respective inequalities for $$$a, b$$$ are both strict, then both $$$a, b$$$ are in the same subtree of $$$c$$$, in which case $$$c$$$ can not be the midpoint. Otherwise, say, for $$$a$$$ we have $$$r \geq dist(v, a) = dist(v, c) + dist(c, a) = dist(v, c) + dist(a, b) / 2$$$, and $$$dist(v, c) \leq r - dist(a, b) / 2$$$ by simple rearrangement.
An immediate corollary of this lemma is this: if a circle $$$C(v, r)$$$ contains a set of proper vertices $$$S$$$, it also contains $$$Cover(S)$$$. We apply the lemma to the midpoint of $$$diam(S)$$$, which is also the center $$$v'$$$ of $$$Cover(S) = C(v', r')$$$, to establish $$$dist(v, v') \leq r - r'$$$. Then, for any vertex $$$w$$$ of $$$Cover(S)$$$ we $$$dist(v, w) \leq dist(v, v') + dist(v', w) \leq r$$$, thus $$$w \in C(v, r)$$$.
Lemma 2. Let $$$S$$$, $$$T$$$ be two non-empty sets of proper vertices. Then the smallest covering circle $$$Cover(S \cup T)$$$ can be found by only knowing $$$Cover(S)$$$ and $$$Cover(T)$$$, by the following rule:
• If $$$Cover(S) \subseteq Cover(T)$$$, then $$$Cover(S \cup T) = Cover(T)$$$.
• If $$$Cover(T) \subseteq Cover(S)$$$, then $$$Cover(S \cup T) = Cover(S)$$$.
• Let none of $$$Cover(S) = (v_S, r_S)$$$ and $$$Cover(T) = (v_T, r_T)$$$ contain the other one. Then, $$$Cover(S \cup T) = C(V, R)$$$ has radius $$$R$$$ equal to $$$(r_S + dist(v_S, v_T) + r_T) / 2$$$ (which is asserted to be an integer), and its center $$$V$$$ is located on the unique path $$$v_S v_T$$$ at distance $$$(r - r_S) / 2$$$ (which is asserted to be an integer) from $$$v_S$$$.
Proof. The first two cases are somewhat easy. Say, in the first case $$$Cover(S \cup T)$$$ has to contain $$$T$$$, then it must contain $$$Cover(T)$$$, and then it contains $$$Cover(S)$$$ by default.
Now, assume the latter case. Any two vertices $$$v \in S$$$, $$$u \in T$$$ satisfy $$$dist(v, u) \leq dist(v, v_S) + dist(v_S, v_T) + dist(v_T, u) \leq r_S + dist(v_S, v_T) + r_T$$$, thus $$$\max dist(v, u) = 2R \leq r_S + dist(v_S, v_T) + r_T$$$. Now, let $$$ab$$$ and $$$cd$$$ be diameter endpoints of $$$S$$$ and $$$T$$$ respectively. Here we can even choose, say, $$$a$$$ outside of $$$Cover(T)$$$, since $$$S \not \subseteq Cover(T)$$$, similarly choose $$$c \not \in Cover(S)$$$. We then must have $$$dist(a, c) = r_S + dist(v_S, v_T) + r_T$$$, thus $$$Cover(S \cup T)$$$ has radius $$$R$$$ at least half of that, and $$$V$$$ is the only possible center candidate.
When one or both of $$$Cover(S)$$$ and $$$Cover(T)$$$ has zero radius, the claim can established pretty much directly.
Okay, we're out of the woods now!
The last lemma allows us to maintain $$$Cover(S) = C(v, r)$$$ for any set $$$S$$$, and further "merge" $$$Cover(S)$$$ and $$$Cover(T)$$$ to obtain $$$Cover(S \cup T)$$$, since all we need is distances between vertices, and finding a vertex on a $$$v_S v_T$$$ at certain distance from $$$v_S$$$. We can also efficiently check if $$$C(v, r)$$$ contains $$$C(v', r')$$$ by verifying $$$dist(v, v') \leq r - r'$$$. All of that can be done with binary lifting and LCA manipulations in $$$O(\log n)$$$ per operation.
At last, let's get back to the divide-and-conquer step. Let us find circles $$$C_1(i) = (v_i, r_i) = Cover(i, \ldots, m)$$$, and $$$C_2(j) = (v_j, r_j) = Cover(m + 1, \ldots, r)$$$ for all relevant $$$i, j$$$. We have $$$C_1(m) = C(m, 0)$$$, and $$$C_1(i)$$$ is a merge of $$$C_1(i + 1)$$$ and $$$C(i, 0)$$$ for any smaller $$$i$$$; similar for $$$C_2(j)$$$.
Now, let's consider the sum $$$\sum_{j = m + 1}^r diam(i, j)$$$ for some $$$i$$$. $$$diam(i, j)$$$ is equal to twice the radius of the merge of $$$C_1(i)$$$ and $$$C_2(j)$$$. Consider increasing $$$j$$$ from $$$m + 1$$$ to $$$r$$$. Since $$$C_2(j) \subseteq C_2(j + 1)$$$ for any $$$j$$$, we have three interesting ranges for $$$j$$$, in order from left to right:
1. $$$j \in [m + 1, t_1(i))$$$: $$$C_2(j) \subseteq C_1(i)$$$;
2. $$$j \in [t_1(i), t_2(i))$$$: $$$C_1(i)$$$ and $$$C_2(j)$$$ are not contained in each other;
3. $$$j \in [t_2(i), r]$$$: $$$C_1(i) \subseteq C_2(j)$$$.
In the first range $$$diam(i, j) = 2r_i$$$, and in the third range $$$diam(i, j) = 2r_j$$$. If we know the boundaries $$$t_1(i)$$$ and $$$t_2(i)$$$, then these can be accounted for with prefix sums on $$$r_j$$$ (and simple multiplcation for $$$r_i$$$).
In the second range we have to sum up $$$r_i + dist(v_i, v_j) + r_j$$$. Again, summands $$$r_i$$$ and $$$r_j$$$ are accounted in the same way. Only $$$\sum_{j \in [t_1(i), t_2(i))} dist(v_i, v_j)$$$ remains.
Instead of computing this directly right now, consider how $$$t_1(i)$$$ and $$$t_2(i)$$$ change as $$$i$$$ decreases from $$$m$$$ to $$$l$$$. We have $$$C_1(i) \subseteq C_1(i - 1)$$$, thus we can conclude that $$$t_1(i)$$$ and $$$t_2(i)$$$ both do not decrease, and the range $$$[t_1(i), t_2(i))$$$ is a "sliding window" with both endpoints moving to the right. Boundaries $$$t_1(i)$$$ and $$$t_2(i)$$$ can thus be maintained with the two-pointers approach, and vertices in the range can be maintained with a queue, where vertices enter from the right and leave from the left.
On top of the queue that supports "push-back" and "pop-front", we'd also like to query "sum of distances from an arbitrary vertex $$$v$$$ to all vertices in the queue". To this end (and only this), we will have to use centroid decomposition in its most basic form. It's fair to say that the margins of this explanation are already too crowded for a proper explanation of this. In short, we can add/remove vertices and query the distance sum in $$$O(\log n)$$$ time. This concludes the divide-and-conquer step description.
With the usual divide-and-conquer time analysis we arrive at an $$$O(n \log^2 n)$$$ solution.