PROBLEM LINK:
Author: Chirag Agarwal
Tester: Ankit Sultana
Editorialist: Ankit Sultana
DIFFICULTY:
Medium-Hard
PREREQUISITES:
Auxiliary Tree
Quick Explanation:
For every query, create the corresponding Auxiliary Tree. Now the Aux Tree will have
at most 2*k - 1 nodes. Now find the cost of each node in the aux tree. Not only that,
if for an edge (u, v) in aux tree, there exist a vertex x on the path between
u and v, then we need to find its cost as well. After we are done with this, we
can push all these values in a vector and just take the first k of them.
Explanation
Firstly I will explain what is an Auxiliary Tree. Say you have a tree T with a vertex
set V, and you have been given a set of vertices S such that S \subseteq T. The set of
vertices of the Auxiliary Tree corresponding to T and S will be a set P such that
P is the minimal set such that S \subseteq P and P is closed under the LCA operation.
What that means is for any two vertices u, v \in P, LCA(u, v) \in P.
As an example, consider the following tree, where the nodes of S have been marked in yellow.
The Auxiliary Tree for the above graph will have the vertex P = {1, 2, 3, 4, 6}
Now how to determine the edge set of the Aux Tree?
For every node u \in P, its parent in the Aux Tree will be an ancestor a of u in the
original tree such that a \in P and depth of a is the highest among all such ancestors.
So the Aux Tree corresponding to the above image is:
To build an Auxiliary tree corresponding to a set S, first add all the vertices in S
to P. Then sort the vertices in the set S
according to their discovery time in a dfs traversal. Now for every pair of adjacent vertices
in the sorted S, add their lca to P. This gives the required minimal superset of P
that is closed under the LCA operation.
To add edges, set the parent of each vertex in the Auxiliary Tree. That can be done using a stack.
Refer to tester’s solution to understand how to do that.
Now coming back to the problem, say the set of vertices we get in a query is S. Now we need to build an
Auxiliary Tree for this set S. We have to find the cost of each vertex u that lies on a path between
some p, q \in S. Note that all such vertices u will be either a vertex in P, or they will lie on a
path between the vertices x, y \in P such that x is a parent of y in the auxiliary graph.
To find the cost for all vertices, we will use partial sums on the Auxiliary Tree.
What we have to do is for every vertex u, we have to add the cost corresponding to all
(p, q) \in P^2 such that u lies on the path between p and q.
Another way to think about this is that we have to add the cost for every path (p, q) to every
vertex u that lies on the path between p and q. For simplicity, assume that there is only one such
path for which we need to add the cost. Assume that the Auxiliary tree is as shown in the image below.
Our strategy will be to first do some precomputation and store some information on each of the
vertices, and then propagate that information across all the nodes to get the cost for each vertex
by performing a single dfs.
For the propagation part, we will propagate values from every vertex to its parent. Assume that in
the given image, we need to add costs on the path from u to v. For that, in the preprocessing
phase, we will add the cost (say x) in both u and v (in the array v_cost).
We will also add x in l = LCA(u, v) in the array lca_cost.
Now as you can see, we can add the cost on all nodes on (u, v) path by returning the value x
from each of the nodes on those path, with the exception of l, which should not return any value.
Finally we will propagate these values using a dfs like the following. Here we are also storing some
information in a vector. This vector stores the cost of a vertex, followed by the
count of vertices with that cost.
#define LL long long
LL dfs(int r, int p, vector<pair<LL, int> > *ptr) {
for(auto &elem: aux[r]) {
if(elem != p) {
v_cost[r] += dfs(elem, r, ptr);
}
}
v_cost[r] -= lca_cost[r]; // Since we have added x twice for every path with lca at r
if(p != 0 && depth[r] > depth[p] + 1) {
ptr->push_back({v_cost[r] - lca_cost[r], depth[r] - depth[p] - 1});
}
ptr->push_back({v_cost[r], 1});
// We don't want to return cost for paths (u, v) such that LCA(u, v) = r
return v_cost[r] - lca_cost[r];
}
Finally, just sort this vector and keep adding cost till the count of taken vertices becomes equal to
X.
Complexity
O(k^2) per query.