Problem Link

Description


You are given an integer array nums and two integers cost1 and cost2. You are allowed to perform either of the following operations any number of times:

  • Choose an index i from nums and increase nums[i] by 1 for a cost of cost1.
  • Choose two different indices i, j, from nums and increase nums[i] and nums[j] by 1 for a cost of cost2.

Return the minimum cost required to make all elements in the array equal.

Since the answer may be very large, return it modulo 109 + 7.

 

Example 1:

Input: nums = [4,1], cost1 = 5, cost2 = 2

Output: 15

Explanation:

The following operations can be performed to make the values equal:

  • Increase nums[1] by 1 for a cost of 5. nums becomes [4,2].
  • Increase nums[1] by 1 for a cost of 5. nums becomes [4,3].
  • Increase nums[1] by 1 for a cost of 5. nums becomes [4,4].

The total cost is 15.

Example 2:

Input: nums = [2,3,3,3,5], cost1 = 2, cost2 = 1

Output: 6

Explanation:

The following operations can be performed to make the values equal:

  • Increase nums[0] and nums[1] by 1 for a cost of 1. nums becomes [3,4,3,3,5].
  • Increase nums[0] and nums[2] by 1 for a cost of 1. nums becomes [4,4,4,3,5].
  • Increase nums[0] and nums[3] by 1 for a cost of 1. nums becomes [5,4,4,4,5].
  • Increase nums[1] and nums[2] by 1 for a cost of 1. nums becomes [5,5,5,4,5].
  • Increase nums[3] by 1 for a cost of 2. nums becomes [5,5,5,5,5].

The total cost is 6.

Example 3:

Input: nums = [3,5,3], cost1 = 1, cost2 = 3

Output: 4

Explanation:

The following operations can be performed to make the values equal:

  • Increase nums[0] by 1 for a cost of 1. nums becomes [4,5,3].
  • Increase nums[0] by 1 for a cost of 1. nums becomes [5,5,3].
  • Increase nums[2] by 1 for a cost of 1. nums becomes [5,5,4].
  • Increase nums[2] by 1 for a cost of 1. nums becomes [5,5,5].

The total cost is 4.

 

Constraints:

  • 1 <= nums.length <= 105
  • 1 <= nums[i] <= 106
  • 1 <= cost1 <= 106
  • 1 <= cost2 <= 106

Solution


Python3

class Solution:
    def minCostToEqualizeArray(self, nums: List[int], cost1: int, cost2: int) -> int:
        N = len(nums)
        MOD = 10 ** 9 + 7
        mmax = max(nums)
 
        # always use cost1 as its cheaper
        if cost1 * 2 <= cost2:
            diff = sum(mmax - x for x in nums)
            return (diff * cost1) % MOD
        
        def cost(target, isOdd):
            if isOdd and target % 2 == 0:
                target += 1
            elif not isOdd and target % 2 == 1:
                target += 1
 
            diff = [target - x for x in nums]
            dominant = max(diff)
            total = sum(diff)
 
            # if maxDiff is the dominant element
            if dominant * 2 > total:
                others = total - dominant
                pairs = min(dominant, others)
                extra = dominant - pairs
 
                return (pairs * cost2 + extra * cost1)
            
            return ((total // 2) * cost2 + (total % 2) * cost1)
        
        res = inf
        for isOdd in [True, False]:
            left, right = mmax, mmax * 3
 
            while left <= right:
                mid1 = left + (right - left) // 3
                mid2 = right - (right - left) // 3
                c1 = cost(mid1, isOdd)
                c2 = cost(mid2, isOdd)
 
                if c1 <= c2:
                    right = mid2 - 1
                else:
                    left = mid1 + 1
            
            res = min(res, cost(left, isOdd))
 
        return res % MOD