Problem statement

https://binarysearch.com/problems/Equalize-List/

Solution

We need to minimize |a1 - x| * c1 + ... + |an - x| * cn by x. We can look at it as |a1 - x| + ... + |a1 - x| + ... Solution of this problem is median of numbers a1...a1 (c1 times), ..., an...an (cn times). We can find median, sorting values and then on each step check if we reached median or not.

Complexity

It is O(n log n) for time and O(n) for space.

Code

class Solution:
    def solve(self, nums, costs):
        if len(nums) <= 1: return 0
        m = (sum(costs) + 1) // 2
        pairs = sorted([a, c] for a, c in zip(nums, costs))
        for i, (a, c) in enumerate(pairs):
            n = min(c, m)
            m -= n
            if m == 0: break

        return sum(abs(a - pairs[i][0]) * c for a, c in pairs)