2448. Minimum Cost to Make Array Equal

Description

You are given two 0-indexed arrays nums and cost consisting each of n positive integers.

You can do the following operation any number of times:

• Increase or decrease any element of the array nums by 1.

The cost of doing one operation on the ith element is cost[i].

Return the minimum total cost such that all the elements of the array nums become equal.

 

Example 1:

Input: nums = [1,3,5,2], cost = [2,3,1,14]
Output: 8
Explanation: We can make all the elements equal to 2 in the following way:
- Increase the 0th element one time. The cost is 2.
- Decrease the 1st element one time. The cost is 3.
- Decrease the 2nd element three times. The cost is 1 + 1 + 1 = 3.
The total cost is 2 + 3 + 3 = 8.
It can be shown that we cannot make the array equal with a smaller cost.

Example 2:

Input: nums = [2,2,2,2,2], cost = [4,2,8,1,3]
Output: 0
Explanation: All the elements are already equal, so no operations are needed.

 

Constraints:

• n == nums.length == cost.length
• 1 <= n <= 105
• 1 <= nums[i], cost[i] <= 106
• Test cases are generated in a way that the output doesn't exceed 2⁵³-1

Solutions

Solution 1: Prefix Sum + Sorting + Enumeration

Let's denote the elements of the array nums as $a_1, a_2, \cdots, a_n$ and the elements of the array cost as $b_1, b_2, \cdots, b_n$. We can assume that $a_1 \leq a_2 \leq \cdots \leq a_n$, i.e., the array nums is sorted in ascending order.

Suppose we change all elements in the array nums to $x$, then the total cost we need is:

$$ \begin{aligned} \sum_{i=1}^{n} \left | a_i-x \right | b_i &= \sum_{i=1}^{k} (x-a_i)b_i + \sum_{i=k+1}^{n} (a_i-x)b_i \ &= x\sum_{i=1}^{k} b_i - \sum_{i=1}^{k} a_ib_i + \sum_{i=k+1}^{n}a_ib_i - x\sum_{i=k+1}^{n}b_i \end{aligned} $$

where $k$ is the number of elements in $a_1, a_2, \cdots, a_n$ that are less than or equal to $x$.

We can use the prefix sum method to calculate $\sum_{i=1}^{k} b_i$ and $\sum_{i=1}^{k} a_ib_i$, as well as $\sum_{i=k+1}^{n}a_ib_i$ and $\sum_{i=k+1}^{n}b_i$.

Then we enumerate $x$, calculate the above four prefix sums, get the total cost mentioned above, and take the minimum value.

The time complexity is $O(n\times \log n)$, where $n$ is the length of the array nums. The main time complexity comes from sorting.

Python3JavaC++GoRust

class Solution:
    def minCost(self, nums: List[int], cost: List[int]) -> int:
        arr = sorted(zip(nums, cost))
        n = len(arr)
        f = [0] * (n + 1)
        g = [0] * (n + 1)
        for i in range(1, n + 1):
            a, b = arr[i - 1]
            f[i] = f[i - 1] + a * b
            g[i] = g[i - 1] + b
        ans = inf
        for i in range(1, n + 1):
            a = arr[i - 1][0]
            l = a * g[i - 1] - f[i - 1]
            r = f[n] - f[i] - a * (g[n] - g[i])
            ans = min(ans, l + r)
        return ans

class Solution {
    public long minCost(int[] nums, int[] cost) {
        int n = nums.length;
        int[][] arr = new int[n][2];
        for (int i = 0; i < n; ++i) {
            arr[i] = new int[] {nums[i], cost[i]};
        }
        Arrays.sort(arr, (a, b) -> a[0] - b[0]);
        long[] f = new long[n + 1];
        long[] g = new long[n + 1];
        for (int i = 1; i <= n; ++i) {
            long a = arr[i - 1][0], b = arr[i - 1][1];
            f[i] = f[i - 1] + a * b;
            g[i] = g[i - 1] + b;
        }
        long ans = Long.MAX_VALUE;
        for (int i = 1; i <= n; ++i) {
            long a = arr[i - 1][0];
            long l = a * g[i - 1] - f[i - 1];
            long r = f[n] - f[i] - a * (g[n] - g[i]);
            ans = Math.min(ans, l + r);
        }
        return ans;
    }
}

using ll = long