PROBLEM LINK:
Author: Dymtro Berezein
Tester: Mugurel Ionut Andreica
Editorialist: Praveen Dhinwa
DIFFICULTY:
simple
PREREQUISITES:
simple math
PROBLEM:
You are given two arrays A and B. For each element B_i, you add B_i to either A_{i - 1} (provided it exists, i.e. i \geq 1), A_i or A_{i + 1} (provided it exists, i.e. if i < n). You have to tell whether you can make all the elements of the array A after all the operations equal or not. If yes, then what is the biggest number that you can achieve?
EXPLANATION:
First thing that you should note that each element of array B is added to some elements of the array A. Also, note that each element is added exactly once.
Due to the above observation, we can notice that sum of all the elements of the array A will be old sum of array A plus sum of array B. Let C denote the final array A we will get after adding elements of B. The sum of all elements of array C will be sum of elements of arrays A and B. Let us denote this sum by S. We want all elements of the array C to be equal, so S should be divisible by n.
Now, each element of the array C will be equal to \frac{S}{n}.
So, we know the final desired value of each element of the array C. It might be possible that this value is not achievable by our operations. If it is possible to achieve this value, then this is the only value that we can get for array C and so by default this will also be biggest value that we can get for each element of C.
Now, for checking whether it is possible to make each element of array equal to S or not. We can use a greedy solution which relies on the fact that for i-th element of array B, we are making changes in array B at quite a few indices which are very near to index i itself.
Let us iterate over array A from left to right and check for each element the set of elements from array B that we can add to get the desired value S.
Let us start from index 1.
- If A_1 > S, then we can never get C_1 = S, as all the elements of array B are positive.
- Otherwise, If we have A_1 < S, then we want to add S - A_1 into array A.
- If B_1 = S - A_1, then we can add B_1 to A_1. Change value of B_1 to 0.
- Otherwise if B_2 = S - A_1, then we can add B_2 to A_1. Change value of B_2 to 0.
- If B_1 + B_2 = S - A_1, then we can add B_1 + B_2 to A_1. Change value of B_1, B_2 both equal to zero.
- Otherwise, we can never make A_1 equal to S.
- Now, we proceed for A_2. We add B_1 into A_2. Note that if the element B_1 of the original array was added to A_1, then we had already changed it to 0. So, it won’t be added twice. Also notice that you can’t omit to add B_1 now, as if it is not added for A_2, it can’t be added for any element A_3 and onwards. After that want to add B_2 and B_3 and want to make A-2 equal to S. This case is very similar to what we did in the above case. We check the conditions correspondingly and proceed.
- Now, after the end, we know that we are able to make elements of A equal to S. We should check the condition that each B_i is indeed used. You can do it by checking whether there is some non-zero element B_i in array B.
Note that we are just doing a single pass over the array and greedily deciding the various conditions to check. This will be an \mathcal{O}(n) solution.