pythonarraysalgorithmtime-complexitybig-o

Minimum number of operations for array of numbers to all equal one number


You have one array of numbers, for example [2, 5, 1]. You have a second array of numbers, for example [8, 4, 3]. For each of the numbers in the second array, how many operations would it take to make the first array all equal that number? You can only increment or decrement by 1 at a time.

To get to 8, it would take (8-2)+(8-5)+(8-1)=16 operations. 
To get to 4, it would take (4-2)+(5-4)+(4-1)=6 operations. 
To get to 3, it would take (3-2)+(5-3)+(3-1)=5 operations. 
So the answer would be [16, 6, 5].

I was able to do this in one line:

answer = [sum(abs(x-y) for x in a1) for y in a2]

But this wasn't fast enough when the arrays can have up to 105 items. How would I be able to do this faster?


Solution

  • Same approach as the others (sort, use bisect to split into the i smaller and the n-i larger values, and look up their sums with precomputed prefix sums), just less code:

    from itertools import accumulate
    from bisect import bisect
    
    def solve(a1, a2):
        a1.sort()
        p1 = 0, *accumulate(a1)
        n = len(a1)
        total = p1[-1]
        return [
            (i*y - p1[i]) + (total-p1[i] - (n-i)*y)
            for y in a2
            for i in [bisect(a1, y)]
        ]
    

    Attempt This Online!

    The expression in the list comprehension could be "optimized" to (2*i-n)*y + total-2*p1[i], but I can't make sense of that and prefer the longer version where I compute the costs for the i smaller values and the n-i larger values separately.