Basic Graph Theory, Union Find Data Structure
Given a Graph G of N nodes and E edges, and the information that each edge weight is
A or A+1, find a spanning tree of weight X. There are Q queries, each with a different
N \le 10^5, Q \le 20, A \le 10^9.
The first step of the solution is to reduce the edge weights to 0 and 1, by subtracting edge weights by A. X can be subtracted by (N - 1) * A.
Observe that the answer is YES iff P \le X \le Q, where P denotes the weight of the minimum
spanning tree and Q denotes the weight of the maximum spanning tree.
To construct a valid solution, find E^* i.e. a set of weight 1 edges that must be present in
any spanning tree of G. If there exists multiple options for E^*, find any one of them.
Then, simulate Kruskal’s again, but this time start with edges in E^*. Then add more weight 1
edges until your spanning tree has weight X. Finally finish off with weight 0 edges.
The complexity of such solution is O(N * alpha(N)) per query, where alpha(N) denotes the inverse ackermann function.
So first step of the solution is to reduce the edge weights by A. This makes life easier,
since now all edge weights are either 0 or 1. We can fix the queries by subtracting
(N - 1) * A from it.
For simplicity, I’ll denote weight of minimum spanning tree as P and weight of maximum spanning tree as Q.
Firstly, it is clear that if X < P or X > Q, then answer is NO.
Take any minimum spanning tree T_1. Also take any maximum spanning tree T_2. Let E^* denote the set of weight one edges which is in T_2 but not in T_1.
Iterate through edges in E^*, and add them one by one to T_1.
After adding each edge, there will be one simple cycle formed.
This cycle will have at-least one edge which is not in T_2 : We have to have atleast one
such edge, if we don’t, there is a cycle in T_2, and that is a contradiction.
Remove any such edge. Repeat this process over all edges in E^*.
At the end you get a spanning tree of weight = Q, and we’ve moved incrementally from P to Q, so we’ve created spanning trees of each weight in [P, Q]. So we have a solution for each X in [P, Q]. This proves that we can find a spanning tree for each X in [P, Q].
Now, one way to create the solution is to simulate the algorithm mentioned in the proof above.
We can do this by binary searching on number of edges from E^* that have been added, since we know that every time we add an edge, we either increase the weight of the spanning tree or it remains the same (it never decreases!). This approach would be N * log(N) * alpha(N) per query, and was intended to time out. There is one more observation which helps us get rid of the log factor.
Let E' be the set of one edges in the minimum spanning tree. Note that we can always construct a maximum spanning tree having all edges in E'. This is crucial, because now we know that whenever we create a cycle in the algorithm described above, we have a 0 weight edge in the cycle. This guarantee of having a 0 edge in the cycle leads us to a rather simple algorithm to compute the spanning tree, which is as follows :-
Find E' (There can be multiple possibilities, any one is fine).
Simulate Kruskal’s again, but this time start with edges in E'. Then, add more weight 1 edges (ensuring that you don’t create a cycle at each step) until you’ve reached the desired weight.
Finally, finish up the tree by adding 0 edges, if needed.
If you manage to not violate any of the constraints, you end up with a solution.
This is N * alpha(N) per query, and it works comfortably within the time limit.
BONUS : There also exists a solution that is O(N) per query, but that was not required for this problem.
O(N * alpha(N) * Q), where alpha(N) denotes the inverse ackermann function.