Minimal cost to set the springs to the same length

August 3, 2024

I am working on the following task:

You are given a sorted collection of springs. Each of them has been produced in certain dimensions, and each stretching, or compression, requires quite a bit of effort.
It takes 1 unit of effort to lengthen or shorten a spring that has not yet been moved by a centimeter. Each successive stretch, or compression, of the same spring requires an additional 1 unit of effort. More precisely, if you want to lengthen a spring by D centimeters, he will exert himself by 1,2,…,D successive units of effort.

Your task is to answer q queries about how much minimal effort it would cost to set the k-shortest springs to the same length.

Constraints: N,Q <= 10⁶, x_i <= 10⁹

Input:
N, Q
x₁, x₂, x₃, …
q₁, q₂, q₃, …

N – number of springs
x_i – length of i-th spring (Note that x_i <= x_i+1)
q_i – how many springs should be set to the same length in i-th question

Output:
Answers to questions – minimal effort it would cost to set the k-shortest springs to the same length. Because answer may be large print it modulo 10^9+7.

My solution attempt:
I discovered that the minimal effort would be achieved when the springs are set to the average length of the springs. However in this approach complexity of algorithm is O(n*q) which is too slow for given constraints.

How can I solve this faster?

In my brutal approach I calculate the average length of the springs and then linearly calculate the cost for every spring. It is too slow for N,Q <= 10⁶ but is enough for n,q <= 10³.

>Solution :

The total effort for a spring with starting length x and target length v is proportional to (x-v)².

For a bunch of N springs, Sum[(x_i-v)²] = Sum[x_i²] – 2v*Sum[x_i] +Nv²

Note that the summations do not depend on the target length at all, so you can just precalculate Sum[x_i²] and Sum[x_i] for all the prefixes of your array in O(n) time.

Then for each query, you can look up the appropriate sums and calculate the minimum cost for any v in O(1) time, giving O(n+q) all together.

As you have already determined, the minimum effort target length is the average spring length Sum[x_i]/N. That’s not always going to be an integer, so just try the 2 integers on either side of it.