I got this question in a company's coding round:
Find the minimum cost to make all the elements of an array equal, when 2 operations can be performed-
- Increment an element by 1 which costs 'a' units
- Decrement an element by 1 which costs 'b' units
Constraints:
Array size <= 10^5
a,b <= 10^5
0 <= element <= 10^6
Auto comment: topic has been updated by DaviddeGea1 (previous revision, new revision, compare).
Fix the point you're trying to make each element equal to, which I'll call a 'convergence point'. Let's start it at $$$0$$$, for which the total cost is just $$$b \cdot \sum arr_i$$$. Sort the array for convenience. Now we'll efficiently iterate through all convergence points. When you increase the convergence point, what happens? The costs of the $$$c_a$$$ elements less than it increase by $$$a$$$, and the costs of the $$$c_b$$$ elements greater than or equal to it decrease by $$$b$$$. Therefore our answer changes by $$$a \cdot c_a - b \cdot c_b$$$. Since the array is now sorted, maintaining $$$c_a$$$ and $$$c_b$$$ is simple. Total complexity: $$$O(10^6 + nlogn)$$$.
A bit more insight. The optimal convergence point, or at least one of them, is always on an array element. Why? Let's imagine it isn't. If it's not between two elements, then it's either greater than or less than all elements, in which case we can move it to the nearest array element and either decrease the cost or keep it as the same. So now, it's between two elements. If, for this point, $$$a \cdot c_a = b \cdot c_b$$$, then it is no less optimal to move the point to the nearest array element. Otherwise, if $$$a \cdot c_a < b \cdot c_b$$$, we can increase the point to decrease the cost, and if $$$a \cdot c_a > b \cdot c_b$$$, we can decrease the point. Thus, for any optimal convergence point not at an array element, it is no less optimal to move it to an array element. And now the complexity is $$$O(nlogn)$$$ and this solution works for arbitrary element bounds, like $$$10^9$$$.
Using this fact and prefix sums from front and back, every element can be checked in $$$O(n)$$$ but sorting is required that makes it $$$O(n \log n)$$$
Hint: Suppose the cost of making all elements equal to $$$T$$$ is $$$C$$$. Then what happens when you move $$$T$$$ one element to the right (i.e. $$$T + 1$$$) or left (i.e. $$$T - 1$$$)?