Taro Logo

Minimum Cost to Make Array Equal

Medium
1 views
2 months ago

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 i<sup>th</sup> element is cost[i]. Return the minimum total cost such that all the elements of the array nums become equal.

For example, consider nums = [1,3,5,2] and cost = [2,3,1,14]. The minimum cost is 8 because we can make all elements equal to 2:

  • Increase the 0th element one time (cost: 2).
  • Decrease the 1st element one time (cost: 3).
  • Decrease the 2nd element three times (cost: 1+1+1 = 3). The total cost is 2 + 3 + 3 = 8.

Another example, if nums = [2,2,2,2,2] and cost = [4,2,8,1,3], the minimum cost is 0 because all elements are already equal.

Sample Answer
class Solution:
    def minCost(self, nums: list[int], cost: list[int]) -> int:
        """Calculates the minimum cost to make all elements in nums equal.

        Args:
            nums (list[int]): A list of positive integers.
            cost (list[int]): A list of positive integers representing the cost of changing each element in nums.

        Returns:
            int: The minimum total cost to make all elements in nums equal.
        """

        # Sort the nums and cost pairs based on nums
        nums_cost = sorted(zip(nums, cost))
        nums = [x[0] for x in nums_cost]
        cost = [x[1] for x in nums_cost]

        # Calculate the prefix sum of the costs
        prefix_cost = [0] * len(cost)
        prefix_cost[0] = cost[0]
        for i in range(1, len(cost)):
            prefix_cost[i] = prefix_cost[i - 1] + cost[i]

        # Binary search to find the optimal target value
        left, right = nums[0], nums[-1]
        min_total_cost = float('inf')
        while left <= right:
            mid = (left + right) // 2
            cost1 = self.calculate_cost(nums, cost, mid)
            cost2 = self.calculate_cost(nums, cost, mid + 1)

            min_total_cost = min(cost1, cost2)

            if cost1 < cost2:
                right = mid - 1
            else:
                left = mid + 1

        return min_total_cost

    def calculate_cost(self, nums: list[int], cost: list[int], target: int) -> int:
        """Calculates the total cost to make all elements in nums equal to target.

        Args:
            nums (list[int]): A list of positive integers.
            cost (list[int]): A list of positive integers representing the cost of changing each element in nums.
            target (int): The target value to make all elements in nums equal to.

        Returns:
            int: The total cost to make all elements in nums equal to target.
        """
        total_cost = 0
        for i in range(len(nums)):
            total_cost += abs(nums[i] - target) * cost[i]
        return total_cost


# Example Usage:
nums1 = [1, 3, 5, 2]
cost1 = [2, 3, 1, 14]
solution = Solution()
result1 = solution.minCost(nums1, cost1)
print(f"Minimum cost for nums1: {result1}")  # Output: 8

nums2 = [2, 2, 2, 2, 2]
cost2 = [4, 2, 8, 1, 3]
result2 = solution.minCost(nums2, cost2)
print(f"Minimum cost for nums2: {result2}")  # Output: 0

nums3 = [1,5]
cost3 = [3,2]
result3 = solution.minCost(nums3, cost3)
print(f"Minimum cost for nums3: {result3}") # Output: 8

nums4 = [335,807,620,508,784,598,35,990,224,59,804,500,761,922,551,679,979,554,809,478,90,394,165,701,85,443,725,950,558,737,105,451,679,862,792,575,738,957,739,503,680,728,893,248,937,778,779,704,692,735,873,877,815,959,376,899,733,665,390,278,975,715,401,200,516,375,784,881,750,772,570,234,856,448,797,472,944,344,285,754,723,876,548,758,920,47,358,451,769,256,147,175,701,962,925,651,946]
cost4 = [89,38,15,52,76,71,75,35,86,58,35,72,78,32,75,67,28,47,75,13,72,51,79,44,45,41,32,56,97,35,62,89,92,72,79,90,60,39,82,71,83,58,91,76,63,37,79,84,33,66,63,70,87,86,95,34,34,63,96,69,40,67,74,47,59,70,79,23,44,91,38,67,84,83,75,26,69,64,49,70,73,92,80,22,79,20,66,58,79,74]
result4 = solution.minCost(nums4, cost4)
print(f"Minimum cost for nums4: {result4}")

Explanation:

The problem asks us to find the minimum cost to make all elements of the array nums equal by either increasing or decreasing its elements. The cost of increasing or decreasing the i-th element by 1 is given by cost[i]. We can solve this problem using the following approach:

  1. Brute Force Approach (Not Implemented):

    • Try all possible values within the range of min(nums) to max(nums) as the target value to which all elements should be equal.
    • For each target value, calculate the total cost by summing the cost of changing each element to the target value.
    • The minimum of all these costs will be the answer.
    • This approach has a time complexity of O(n * (max(nums) - min(nums))), which can be very slow for large ranges.
  2. Optimal Approach (Implemented):

    • The optimal target value to which all elements should be equal lies between the minimum and maximum values in nums. We can use binary search to find this optimal value.
    • The cost function is convex, meaning it has a single minimum. We exploit this property to efficiently find the minimum cost.
    • We start by sorting the nums and cost pairs based on nums. This allows us to keep track of the cost associated with each number.
    • We then use binary search between the minimum and maximum values of nums. In each iteration, we calculate the cost for two adjacent target values, mid and mid + 1. If the cost for mid is less than the cost for mid + 1, it means the optimal target value lies to the left of mid, so we update the right boundary. Otherwise, it lies to the right, so we update the left boundary.
    • We continue this process until the left and right boundaries converge. The minimum cost found during the binary search is the answer.

Big(O) Run-time Analysis:

  • Sorting: O(n log n), where n is the number of elements in nums. This is due to the sorted() function.
  • Binary Search: O(log(max(nums) - min(nums))). In each step, we are halving the search space.
  • Cost Calculation: O(n) for each target value in the binary search. We iterate through the nums array once to calculate the cost.
  • Overall: The dominant factor is the sorting step and the binary search with cost calculation. Therefore, the time complexity is O(n log n + n * log(max(nums) - min(nums))). In the worst case, where the range (max(nums) - min(nums)) is very large, this could approach O(n log n + n log M), where M is the difference between max and min.

Big(O) Space Usage Analysis:

  • Sorted List: O(n) to store the sorted list of nums and cost pairs.
  • Prefix Cost: O(n) for the prefix cost array, if used. However, the provided solution does not explicitly use the prefix sum array. In this case, it would be O(1).
  • Overall: The space complexity is primarily determined by the sorted list and is O(n).

Edge Cases:

  1. Empty Input: If nums is empty, we should return 0. However, the problem states that 1 <= n <= 10^5, so we don't need to handle this case.
  2. All Elements are Equal: If all elements in nums are already equal, the cost should be 0. The given solution handles this case correctly.
  3. Large Input Values: The problem statement mentions that the output doesn't exceed 2^53 - 1. We should ensure that our calculations don't overflow. The code uses long to prevent integer overflow during cost calculation.
  4. Negative Input Values: The problem statement specifies positive integers, so negative values are not an issue. If negative values were allowed, the binary search would still work correctly since it searches within the range of input values, and the cost calculation uses absolute differences.
  5. Cost values of 0: The costs are given to be positive integers, so there is no need to handle cases where the cost is zero. However, the algorithm should still perform correctly and yield optimal results.