Skip to content

3139. Minimum Cost to Equalize Array 👍

  • Time: $O(\max(\texttt{nums}))$
  • Space: $O(1)$
 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
class Solution {
 public:
  int minCostToEqualizeArray(vector<int>& nums, int cost1, int cost2) {
    constexpr int kMod = 1'000'000'007;
    const int n = nums.size();
    const int minNum = ranges::min(nums);
    const int maxNum = ranges::max(nums);
    const long sum = accumulate(nums.begin(), nums.end(), 0L);
    long ans = LONG_MAX;

    if (cost1 * 2 <= cost2 || n < 3) {
      const long totalGap = static_cast<long>(maxNum) * n - sum;
      return (cost1 * totalGap) % kMod;
    }

    for (int target = maxNum; target < 2 * maxNum; ++target) {
      const int maxGap = target - minNum;
      const long totalGap = static_cast<long>(target) * n - sum;
      const long pairs = min(totalGap / 2, totalGap - maxGap);
      ans = min(ans, cost1 * (totalGap - 2 * pairs) + cost2 * pairs);
    }

    return ans % kMod;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class Solution {
  public int minCostToEqualizeArray(int[] nums, int cost1, int cost2) {
    final int MOD = 1_000_000_007;
    final int n = nums.length;
    final int minNum = Arrays.stream(nums).min().getAsInt();
    final int maxNum = Arrays.stream(nums).max().getAsInt();
    final long sum = Arrays.stream(nums).asLongStream().sum();
    long ans = Long.MAX_VALUE;

    if (cost1 * 2 <= cost2 || n < 3) {
      final long totalGap = 1L * maxNum * n - sum;
      return (int) ((cost1 * totalGap) % MOD);
    }

    for (int target = maxNum; target < 2 * maxNum; ++target) {
      final int maxGap = target - minNum;
      final long totalGap = 1L * target * n - sum;
      final long pairs = Math.min(totalGap / 2, totalGap - maxGap);
      ans = Math.min(ans, cost1 * (totalGap - 2 * pairs) + cost2 * pairs);
    }

    return (int) (ans % MOD);
  }
}

 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
class Solution:
  def minCostToEqualizeArray(
      self,
      nums: list[int],
      cost1: int,
      cost2: int,
  ) -> int:
    MOD = 1_000_000_007
    n = len(nums)
    minNum = min(nums)
    maxNum = max(nums)
    summ = sum(nums)

    if cost1 * 2 <= cost2 or n < 3:
      totalGap = maxNum * n - summ
      return (cost1 * totalGap) % MOD

    def getMinCost(target: int) -> int:
      """Returns the minimum cost to make all numbers equal to `target`."""
      maxGap = target - minNum
      totalGap = target * n - summ
      # Pair one shallowest number with one non-shallowest number, so the worst
      # case is that we have `totalGap - maxGap` non-shallowest numbers to pair.
      pairs = min(totalGap // 2, totalGap - maxGap)
      return cost1 * (totalGap - 2 * pairs) + cost2 * pairs

    return min(getMinCost(target)
               for target in range(maxNum, 2 * maxNum)) % MOD