PROBLEM LINK:
Author: Vipul Sharma
Tester: Prateek Gupta
Editorialist: Oleksandr Kulkov
DIFFICULTY:
MEDIUM
PREREQUISITES:
None
PROBLEM:
You are given three arrays A, B and C. You have to sum up (A_i+B_j)(B_j+C_k) for all i, j, k such that A_i \leq B_j \geq C_k.
QUICK EXPLANATION:
Over all B_j's, sum up \left(\sum\limits_{A_i \leq B_j}(A_i+B_j)\right)\times\left(\sum\limits_{C_k \geq B_j}(B_j+C_k)\right).
EXPLANATION:
Let’s bruteforce B_j. If we fix some j, then we can choose A_i \leq B_j and C_k \geq B_j independently. Let’s sort A in ascending order and C in descending order. Then suitable A_i as well as suitable C_k will form prefix of A and C correspondingly which size can be found by binary search. Let the sizes of these prefixes will be k_1 and k_2. Then from triples with this B_j we will have contribution of \sum\limits_{i=1}^{k_1}\sum\limits_{k=1}^{k_2}(A_i+B_j)(B_j+C_k)=\sum\limits_{i=1}^{k_1}(A_i+B_j)\left(k_2 B_j+\sum\limits_{k=1}^{k_2}C_k\right)=\left(k_1 B_j+\sum\limits_{i=1}^{k_1}A_i\right)\times \left(k_2 B_j+\sum\limits_{k=1}^{k_2}C_k\right) which can be quickly calculated using prefix sums. Thus we have O(n \log n) solution.
We can optimize this solution to \mathcal{O}(n) by making the following observation. Let us sort the B in increasing order. The values of k_1 and k_2 can only increase when we go from left to right in the array B. This technique’s standard name is “two pointers” approach.
AUTHOR’S AND TESTER’S SOLUTIONS:
Author’s solution will be updated soon.
Tester’s solution will be updated soon.