Wheels
Contest: Kotlin Heroes: Episode 1
Tags: *special, binary search, greedy

Polycarp has $$$n$$$ wheels and a car with $$$m$$$ slots for wheels. The initial pressure in the $$$i$$$-th wheel is $$$a_i$$$.

Polycarp's goal is to take exactly $$$m$$$ wheels among the given $$$n$$$ wheels and equalize the pressure in them (then he can put these wheels in a car and use it for driving). In one minute he can decrease or increase the pressure in any (single) wheel by $$$1$$$. He can increase the pressure no more than $$$k$$$ times in total because it is hard to pump up wheels.

Help Polycarp and say what is the minimum number of minutes he needs to spend to equalize the pressure of at least $$$m$$$ wheels among the given $$$n$$$ wheels.

Input format:
The first line of the input contains three integers $$$n, m$$$ and $$$k$$$ ($$$1 \le m \le n \le 2 \cdot 10^5, 0 \le k \le 10^9$$$) — the number of wheels, the number of slots for wheels in a car and the number of times Polycarp can increase by $$$1$$$ the pressure in a wheel.

The second line of the input contains $$$n$$$ integers $$$a_1, a_2, \dots, a_n$$$ ($$$1 \le a_i \le 10^9$$$), where $$$a_i$$$ is the pressure in the $$$i$$$-th wheel.

Output format:
Print one integer — the minimum number of minutes Polycarp needs to spend to equalize the pressure in at least $$$m$$$ wheels among the given $$$n$$$ wheels.

Constraints:
Time limit: 3.0 s
Memory limit: 256.0 MB

Example input:
6 6 7
6 15 16 20 1 5

Example output:
39

Example input:
6 3 1
4 8 15 16 23 42

Example output:
8

Editorial:
It is obvious that it is always better to try to equalize the subsegment of the sorted array. So let's sort the array and iterate over all $$$n-m+1$$$ of its subsegments.
Let $$$sum(l, r) = \sum\limits_{j=l}^{r} a_j$$$.
For the segment $$$[i; i + m - 1]$$$ we have to find the minimum number of moves needed to equalize this segment. In the best case (if $$$k=\infty$$$) it can be proven that the better way to equalize all elements is to move them to the middle element of the segment (or if $$$m$$$ is even then we can choose either the left middle or the right middle element). How to calculate answer for this case? Let $$$idx$$$ be the middle of the segment (in other words, $$$i + \lfloor\frac{m}{2}\rfloor$$$), $$$val$$$ be the middle element (in other words, $$$a_{idx}$$$) and $$$cnt$$$ be $$$idx - i + 1$$$. Then the number of moves we need is $$$res = (cnt \cdot val - sum(i, idx)) + (sum(idx + 1, i + m - 1) - (m - cnt) \cdot val)$$$. This formula gives the number of operations needed to equalize the whole segment to the value $$$val$$$. The left part of the formula increases the part to the left of $$$val$$$ and the rightmost part decreases the part to the right of $$$val$$$.
This is almost the complete solution! What should we do with $$$k$$$? Let's find the rightmost position $$$idx$$$ that belongs to the segment and we can increase all elements in segment $$$[i; idx]$$$ to the value $$$a_{idx}$$$. We can do it using binary search and prefix sums (or two pointers). Then consider two cases:
1. $$$idx < i + \lfloor\frac{m}{2}\rfloor$$$;
2. $$$idx \ge i + \lfloor\frac{m}{2}\rfloor$$$.
In the first case $$$idx$$$ does not change, $$$cnt = i + idx - 1$$$ and $$$val = a_{idx} + \lfloor\frac{k - (cnt \cdot a_{idx} - sum(i, idx))}{cnt}\rfloor$$$. Then we can substitute these variables to the formula calculating $$$res$$$ and update the answer with the result. What does this scary formula ($$$\lfloor\frac{k - (cnt \cdot a_{idx} - sum(i, idx))}{cnt}\rfloor$$$) mean? It means that we can increase all values on the segment $$$[i; idx]$$$ to the value $$$a_{idx}$$$ but there are some increasing opreations remain and we should use them. And because $$$idx$$$ is the rightmost such position, it means that we cannot increase all values to $$$a_{idx + 1}$$$ but we can increase the whole segment $$$[i; idx]$$$ by some value possible greater than $$$a_{idx}$$$ (and it will not increase the result for this segment).
In the second case $$$idx = i + \lfloor\frac{m}{2}\rfloor$$$, $$$cnt = idx - i + 1$$$ and $$$val = a_{idx}$$$, and we can substitute these variables into the formula calculating $$$res$$$ and update the answer with the result.