PROBLEM LINK
Author: Hanlin Ren
Tester: Jakub Safin
Editorialist: Jakub Safin
DIFFICULTY
HARD
PREREQUISITIES
centroid decomposition, Fenwick trees, preorder sorting, lowest common ancestor
PROBLEM
You’re given two trees on the same set of N vertices. Let d_i(u,v) be the distance between vertices u and v in tree number i. Compute the sum of \mathrm{min}(d_1,d_2) for all pairs of vertices.
QUICK EXPLANATION
Use centroid decomposition on tree 1, compressing tree 2 when recursing into subtrees. To compute the sum with u,v on different sides from the centroid, use another centroid decomposition on tree 2. To compute sums of the form \mathrm{min}\left(a(u)+a(v),b(u)+b(v)\right), use 3 Fenwick trees.
EXPLANATION
For an O(N^2) solution, it’s enough to calculate all distances in both trees by running a DFS from each vertex and computing the answer by brute force. This should be enough to solve the first subtask, but fails miserably on subtasks 2-4.
In subtask 2, we just need the sum of distances in one tree, which is easy: for each edge of length l, if removing it splits the tree into subtrees with n and N-n vertices, add n(N-n)l to the answer.
For the 4th subtask, there’s an O(N\log^3 N) solution using a well-known tree technique, centroid decomposition - twice. (Using it just once solves the 3rd subtask.)
First of all, let’s extend the problem by marking some vertices as dead (in both trees) and computing only \sum\mathrm{min}(d_1,d_2) for pairs of living vertices. Initially, all vertices are living.
Now let’s find the centroid c of T_1, root T_1 at c and compute all distances to the root in T_1; let the distance to vertex v be a(v). We need to compute \sum\min(d_2(u,v),a(u)+a(v)) for all pairs of living vertices that aren’t in the same subtrees. We can do that by first computing the sum for all pairs of living vertices and then subtracting the sums for all subtrees; that can be done in the same way after recursing into subtrees.
Afterwards, we need to compute the sums for all subtrees of T_1 created after removing c from it. There’s a small problem here: we can’t take subtrees after removing c from T_2 and can’t take full copies of T_2 either - that’s too slow. For each subtree, we need to generate a smaller tree from T_2 that contains all living vertices in it and maybe some more dead vertices so that all paths would remain the same as in T_2.
We can sort the living vertices in preorder and take all LCAs of adjacent living vertices to be the set of all additional, possibly dead vertices. This works because of the following lemma:
If u < v < w in the preorder numbering, then \mathrm{LCA}(u,w)=\mathrm{LCA}(u,v) or \mathrm{LCA}(u,w)=\mathrm{LCA}(v,w).
Proof: Denote \mathrm{LCA}(u,v,w)=l. Let’s prove l=\mathrm{LCA}(u,w). If it doesn’t hold, the former is clearly an ancestor of the latter, with v in a different subtree of l than u,w or equal to l. In both cases, the subtree containing u,w corresponds to an interval in preorder that doesn’t contain v, which is a contradiction with u < v < w. Now let’s prove that l is equal to \mathrm{LCA}(u,v) or \mathrm{LCA}(v,w), by contradiction. If \mathrm{LCA}(u,v) \neq l, then u,v \neq l and are in the same subtree of l; similarly for \mathrm{LCA}(v,w) \neq l. However, u,v,w are then in the same subtree of l, which is a contradiction. QED
It turns out that this way, any subtree containing N living vertices only needs to contain at most N dead vertices, everything else is contracted into edges. Therefore, it doesn’t affect time complexity.
Now we have vertices, but need to find the edges between them too. If we sort the vertices in preorder again, then we can process them one by one, keep a stack like in a DFS, remove vertices on the stack that aren’t ancestors of the current vertex and add an edge between the last vertex in the stack and the current vertex. This is similar to how the preorder numbering of vertices can be done using DFS and takes linear time.
Sum with distances from c
We still need to find out how to compute \sum\min(d_2(u,v),a(u)+a(v)). Let’s use another centroid decomposition, this time on T_2. We don’t need to split T_1 now since it doesn’t enter into the sum directly. With a given centroid c_2 and distances b(v) to v, we need to compute \sum\min(b(u)+b(v),a(u)+a(v)). There’s no need to worry about excluding pairs of vertices in the same subtrees of c_2, we can simply take u to be all vertices in already processed subtrees and sum over all v in the current subtree.
Let’s deal with the sum this way: for a given v, we’re summing up:
- b(u)+b(v) over u with b(u)-a(u) \le a(v)-b(v),
- a(u)+a(v) over u with b(u)-a(u) > a(v)-b(v).
If there are n such vertices, it’s equal to n b(v) + \sum b(u) for the first type and n a(v) + \sum a(u) for the second type. Such prefix/suffix sums can be computed easily with segment trees or Fenwick trees in O(\log{D}), where D is the maximum distance in one tree.
However, D can be very large here, so we need to compress all b(v)-a(v) and a(v)-b(v) into a smaller range. The best solution is putting them in an array and sorting; using STL map gives the same complexity, but with a larger constant factor. This way, we only need O(N) memory.
Since processing one subtree with n vertices takes O(n\log{n}) and each vertex appears in O(\log{N}) vertices, this sum can be computed in O(N\log^2{N}). The first centroid decomposition adds another \log{N} to the complexity. Sorting and computing LCAs takes O(\log{N}) per vertex, so that’s just O(N\log^2{N}) in total - the whole algorithm takes O(N\log^3{N}) time and O(N\log{N}) memory.
