PROBLEM LINK:
Author: Shashwat Chandra
Tester: Shashwat Chandra, Taranpreet Singh, Bharat Goyal
Editorialist: Shashwat Chandra
DIFFICULTY:
MEDIUM-HARD
PREREQUISITES:
Segment Tree, Binary Search, Trees, Persistence
EXPLANATION:
Let’s fix the proton at some node A, it is clear from the statement that the E electrons can only be placed at nodes which are either an ancestor of A or a descendant (lie in A's subtree).
Let us also fix the exact number of electrons which lie above and below A to be P and Q respectively. Here, P+Q = E.
OBSERVATION 1
Now consider the minimum and maximum possible values of S_1, let them be L_1 and R_1 respectively.
Similarly consider the minimum and maximum possible values of S_2, let them be L_2 and R_2 respectively.
We claim that every possible value of L_1 \leq S_1 \leq R_1 is attainable and similarly every possible value of L_2 \leq S_2 \leq R_2 is attainable.
The proof is given below.
Proof
Consider an array which contains distances from A to all nodes in its subtree(or ancestors), let us sort it.
The minimum possible sum M_1 is attained by choosing the first K values and the maximum sum M_2 by choosing the last K.
Initially we select the first K thereby getting the minimum possible sum.
Now suppose, currently the sum of all selected positions is X < M_2, we provide a method of going to X+1 :
Step 1: Find any index i such that i is selected but i+1 is not.
Step 2: Remove i from the selected set of indices and add i+1.
Step 3: If the updated sum is X+1, end. Otherwise, go to Step 1.
Note that this works because whenever we delete some index i and add i+1, we increase the sum by 0 or 1.
This observation enables us to solve the problem in O(N^2 *log_2(N)).
The solution is as follows:
- We calculate the answer for each node while doing DFS on the tree.
- For each node, have the distances to all nodes in it subtree stored and sorted.
- Compute prefix sums for this array.
- Iterate over every possible value of P \geq X_A and Q \geq Y_A and calculate the values of L_1,R_1,L_2,R_2. Also, shift [L_1,R_1] by Z_A.
- If [L_1,R_1] intersects [L_2,R_2] , the answer for this P,Q is 0.
- Otherwise, the answer for this P,Q is min(|L_1-R_2|,|L_2-R_1|)
- The final answer for A is the minimum of all intermediate answers obtained.
Note that for each A we spend O(N*log_2(N)) time for computing the distance array and O(N) time for calculating the answer for all possible P and Q.
From now on, whenever we mention L_1 and R_1, we assume they have been shifted by Z_A already.
OBSERVATION 2
As we decrease P (thereby increasing Q) , [L_1,R_1] keeps shifting to the left and similarly [L_2,R_2] keeps shifting to the right. This is quite intuitive hence the proof is omitted.
What we can gather from this observation is that, first both the ranges come towards each other and then after some time they pass through and move away.
We can further say that the minimum value will be attained either when they overlap (here, the answer is 0). In cases where they never overlap, there are at most 2 values of P,Q which can give us the optimal answer, the first when [L_1,R_1] is just ahead of [L_2,R_2] and the next when [L_2,R_2] is just ahead of [L_1,R_1].
Note that these 2 cases are consecutive as in, the first case is when P = a,Q = b while the second case is when P = a-1,Q = b+1.
We can now easily find the values of a and b by binary searching on a.
We just need to check the values for L_1,R_1,L_2,R_2 at P = a and P = a-1.
Refer to the codes provided at the end for more details.
We will also need to adjust the low and high values in the binary search so as to satisfy the X_A and Y_A constraints.
Reduction to a well-known data structure problem
Now all we need to do is find a way to compute L_1,R_1,L_2,R_2 quickly, say in time T(n) then we can solve the problem in O(N*log_2(N)*T(N)).
Actually computing L_1,R_1 can be done in O(1) since its just a sum of P consecutive numbers which can be computed using the formula for A.P.
For L_2,R_2 , we need to compute the sum of Q largest/smallest numbers in the subtree of A.
We can do this by first flattening the tree i.e creating the DFS array. After this the problem is to find the sum of Q largest numbers in a range [L,R]. (The sum of the Q smallest values can be computed in a similar fashion).
This is a well-known problem and can be solved by a persistent segment tree.
The pseudocode is given below.
Code
Query(L,R,Left,Sum):
if L == R:
return Sum+Left*L;
if #right_child >= Left:
return Query(Mid+1,R,Left,Sum)
return Query(L,Mid,Left-#right_child,Sum+seg(right_child))
Here L,R denote the values covered by the current node, #right_child is the number of values present in right child of the current node and seg(right_child) is the sum of those values.
Note we must get seg(right_child) and #right_child by using the values from the 2 versions of the segment tree which tell us about the range [L,R] namely the R^{th} and (L-1)^{th} versions.
Note that this approach takes O(N*log_2(N)) preprocessing and T(N) = log_2(N), thus the overall complexity is N*log_2(N)^2).