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
fromnums
and increasenums[i]
by1
for a cost ofcost1
. - Choose two different indices
i
,j
, fromnums
and increasenums[i]
andnums[j]
by1
for a cost ofcost2
.
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]
andnums[1]
by 1 for a cost of 1.nums
becomes[3,4,3,3,5]
. - Increase
nums[0]
andnums[2]
by 1 for a cost of 1.nums
becomes[4,4,4,3,5]
. - Increase
nums[0]
andnums[3]
by 1 for a cost of 1.nums
becomes[5,4,4,4,5]
. - Increase
nums[1]
andnums[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