Problem

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.

Examples

Example 1:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

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:

1
2
3
4
5

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.

Solution

Method 1 - Sorting

To solve the problem where we need to equalize all elements in the nums array with the minimum total cost, we need to leverage the properties of weighted medians.

Weighted Median Explanation

  • If we want to minimize the cost to make all elements equal to a common value x, the optimal x is the weighted median of the elements, where the weights are given by the cost array.
  • A weighted median is a value that minimizes the sum of weighted absolute differences: (\sum (|num_i - x|) \times cost_i).

Here is the approach:

  1. Sort: Create tuples of (nums[i], cost[i]) and sort them based on the values of nums.
  2. Calculate the Total Weight: The total weight is the sum of all elements in the cost array.
  3. Find the Weighted Median: Traverse the sorted array and find the point where the cumulative sum of weights is at least half of the total weight.
  4. Calculate the Total Cost: Using the weighted median found, compute the total cost to make all elements equal to this value.

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37

public class Solution {
    public long minCost(int[] nums, int[] cost) {
        int n = nums.length;
        long[][] pairs = new long[n][2];
        for (int i = 0; i < n; i++) {
            pairs[i][0] = nums[i];
            pairs[i][1] = cost[i];
        }

        // Sort pairs by nums values
        Arrays.sort(pairs, (a, b) -> Long.compare(a[0], b[0]));

        // Find the weighted median
        long totalCost = 0;
        for (int c : cost) {
            totalCost += c;
        }

        long cumulativeCost = 0;
        long median = 0;
        for (int i = 0; i < n; i++) {
            cumulativeCost += pairs[i][1];
            if (cumulativeCost >= (totalCost + 1) / 2) {
                median = pairs[i][0];
                break;
            }
        }

        // Calculate the total cost to make all elements equal to the median
        long minTotalCost = 0;
        for (int i = 0; i < n; i++) {
            minTotalCost += Math.abs(nums[i] - median) * cost[i];
        }

        return minTotalCost;
    }
}
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

class Solution:
    def minCost(self, nums: List[int], cost: List[int]) -> int:
        combined = sorted(zip(nums, cost))

        # Calculate the total weight
        total_weight = sum(cost)

        # Find the weighted median
        cumulative_weight = 0
        median = 0
        for num, c in combined:
            cumulative_weight += c
            if cumulative_weight >= (total_weight + 1) // 2:
                median = num
                break

        # Calculate the total cost to make all elements equal to the median
        min_total_cost = 0
        for num, c in zip(nums, cost):
            min_total_cost += abs(num - median) * c

        return min_total_cost

Complexity

  • ⏰ Time complexity: O(n log n)
    • Sorting: Sorting the combined nums and cost tuples takes O(n log n).
    • Weighted Median Calculation: This traversal is O(n).
    • Total Cost Calculation: This traversal is also O(n).
    • Thus, the overall time complexity is dominated by the sorting step, resulting in O(n log n).
  • 🧺 Space complexity: O(n), for storing the tuples of nums and cost.