Almost Same Distance
Contest: Codeforces Global Round 6
Difficulty rating: 2900
Tags: dfs and similar, graphs

Let $$$G$$$ be a simple graph. Let $$$W$$$ be a non-empty subset of vertices. Then $$$W$$$ is almost-$$$k$$$-uniform if for each pair of distinct vertices $$$u,v \in W$$$ the distance between $$$u$$$ and $$$v$$$ is either $$$k$$$ or $$$k+1$$$.

You are given a tree on $$$n$$$ vertices. For each $$$i$$$ between $$$1$$$ and $$$n$$$, find the maximum size of an almost-$$$i$$$-uniform set.

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

Then $$$n-1$$$ lines follows, the $$$i$$$-th of which consisting of two space separated integers $$$u_i$$$, $$$v_i$$$ ($$$1 \leq u_i, v_i \leq n$$$) meaning that there is an edge between vertices $$$u_i$$$ and $$$v_i$$$.

It is guaranteed that the given graph is tree.

Output format:
Output a single line containing $$$n$$$ space separated integers $$$a_i$$$, where $$$a_i$$$ is the maximum size of an almost-$$$i$$$-uniform set.

Constraints:
Time limit: 5.0 s
Memory limit: 256.0 MB

Note:
Consider the first example.

- The only maximum almost-$$$1$$$-uniform set is $$$\{1, 2, 3, 4\}$$$.
- One of the maximum almost-$$$2$$$-uniform sets is or $$$\{2, 3, 5\}$$$, another one is $$$\{2, 3, 4\}$$$.
- A maximum almost-$$$3$$$-uniform set is any pair of vertices on distance $$$3$$$.
- Any single vertex is an almost-$$$k$$$-uniform set for $$$k \geq 1$$$.

In the second sample there is an almost-$$$2$$$-uniform set of size $$$4$$$, and that is $$$\{2, 3, 5, 6\}$$$.

Example input:
5
1 2
1 3
1 4
4 5

Example output:
4 3 2 1 1

Example input:
6
1 2
1 3
1 4
4 5
4 6

Example output:
4 4 2 1 1 1

Editorial:
There are three cases of how an almost-$$$k$$$-uniform set looks like depending on the value of $$$k$$$.
• The first case is when $$$k = 1$$$. In this case, any maximal almost-$$$1$$$-uniform set is a vertex plus all of its neighbours. We can simply check each vertex and find the one with the highest degree.
• The second case is when $$$k$$$ is odd and greater than one: $$$k = 2l + 1$$$ for $$$l \geq 1$$$. Then any almost-$$$k$$$-uniform set looks as follows. There is a middle vertex $$$v$$$, and every $$$w \in W$$$ has distance of $$$l$$$ or $$$l+1$$$ from $$$v$$$. Additionally, at most one of the vertices is in distance $$$l$$$, and each $$$w$$$ is in a different subtree of $$$w$$$.
• The third case is when $$$k$$$ is even: $$$k = 2l$$$ for $$$l \geq 1$$$. Then any almost-$$$k$$$-uniform set can be constructed in one of two ways. The first is similar to the previous case, except that all vertices have to be at distance of exactly $$$l$$$ from the middle vertex. The second construction begins by selecting a middle edge $$$\{u, v\}$$$, removing it from the tree and then finding a set of vertices $$$W$$$ such that each $$$w \in W$$$ is in the distance of $$$l$$$ from either $$$u$$$ or $$$v$$$ (depending in which subtree it is) and each $$$w$$$ is in a different subtree of $$$u$$$ or $$$v$$$.
Now we need to figure out how to make these constructions efficiently. We begin by calculating $$$d_v(u)$$$ – the depth of a subtree $$$u$$$ when the tree is rooted at $$$v$$$. We compute this for each edge $$$\{u, v\}$$$ by using the standard technique of two DFS traversals. Furthermore we use the fact that the answer is monotone with respect to parity: $$$Ans[i+2] \leq Ans[i]$$$ for all $$$i$$$ – we only calculate the answer in some "interesting" points and then preserve this inequality in the end by taking suffix maximums.
For the odd case, consider a middle vertex $$$v$$$. Sort the depths of subtrees of neighbours of $$$v$$$. Then, $$$Ans[2*l+1] \geq x$$$ if there are at $$$x-1$$$ subtrees of depth at least $$$l+1$$$ and one additional subtree of depth $$$l$$$ or more. By sorting the depths, we can process each middle vertex in $$$\mathcal O(\mathrm{deg}_v \log \mathrm{deg}(v))$$$. In total, we have $$$\mathcal O(n \log n)$$$ for this step.
For the even case there are two options. The first of these is very similar to the above construction, we just look for $$$x$$$ subtrees of depth at least $$$l$$$. The second option is more involved. Consider a middle edge $$$\{u,v\}$$$. Let $$$x$$$ be the number of $$$d_u(w) \geq l$$$ for $$$w \neq v$$$ and $$$y$$$ be the number of $$$d_v(w) \geq l$$$ for $$$w \neq u$$$. Then we can conclude that $$$Ans[2*l] \geq x + y$$$.
However, we cannot directly calculate the above quantity for each edge, as the processing of each vertex would take time quadratic in its degree. We will do a series of steps to improve upon this.
The first optimization is that in the process of finding $$$x$$$ and $$$y$$$ above, we also consider $$$d_v(u)$$$ and $$$d_u(v)$$$, but then subtract $$$2$$$ from the answer. Why can we do that? First see that this case is only important if both $$$x \geq 1$$$ and $$$y \geq 1$$$ – otherwise this is the middle vertex case that is already processed. This means that $$$\max d_v(w) \geq l$$$ and hence $$$d_u(v) \geq l+1 \geq l$$$. This means that adding $$$d_u(v)$$$ to the set increases $$$x$$$ by one. The same argument shows that $$$y$$$ increases by $$$1$$$.
How does this modification help us? There are two ways of proceeding now.
• We can fix $$$v$$$ as one of the endpoints of the middle edge and merge the subtree depths from individual subtrees by taking the maximum and process all neighbours of $$$u$$$ at the same time. This alone still does not change the complexity, because we still process each vertex once for each of its neighbours, and this processing takes time at least linear in the degree. Now comes the final step.Perform a centroid decomposition of the tree. When processing $$$v$$$ as the fixed end of the middle vertex, consider all $$$d_v(w)$$$ for the almost-$$$k$$$-universal set, but as the other endpoint of the middle edge we consider only the vertices in the tree where $$$v$$$ is the centroid. This way, each middle edge is processed exactly once. Each individual vertex is processed once when it is centroid and then once within each centroid tree it belongs to. Since the depth of centroid decomposition is $$$\mathcal O(\log n)$$$, we only process each vertex $$$\mathcal O(\log n)$$$ times. A single processing every vertex costs $$$\mathcal O(\mathrm{deg}_v \log \mathrm{deg}_v)$$$ time (because of sorting the degrees). Thus the total running time is $$$\mathcal O(n \log^2 n)$$$, which is the most expensive part of the algorithm. We can also replace the sort with parallel counting sort, and this makes the time $$$\mathcal O(n \log n)$$$, but that doesn't run faster in practice.
• For each vertex $$$u$$$, store the histogram of $$$d_u(v)$$$ in a map. Now, consider each edge separately, and merge calculate the answer naively, using a linear pass through $$$d_v(u)$$$ and $$$d_u(v)$$$. Why does it work fast? The sum of all $$$d_u(v)$$$ for a fixed $$$u$$$ equals to $$$n-1$$$, so there are at most $$$\mathcal O(\sqrt n)$$$ different values in the histogram, which makes the solution $$$\mathcal O(n \sqrt n)$$$ in total, and quite fast in practice.