Search code examples
algorithmmath

Minimal cost to set the springs to the same length

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 <= 106, xi <= 109

Input: N, Q
x1, x2, x3, ...
q1, q2, q3, ...

N - number of springs
xi - length of i-th spring (Note that xi <= xi+1)
qi - 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 <= 106 but is enough for n,q <= 103.

Solution

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

    For a bunch of N springs, Sum[(xi-v)2] = Sum[xi2] - 2v*Sum[xi] +Nv2

    Note that the summations do not depend on the target length at all, so you can just precalculate Sum[xi2] and Sum[xi] 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[xi]/N. That's not always going to be an integer, so just try the 2 integers on either side of it.