Problem of The Day: Minimum Cost to Make Array Equal

5 min read 1180 words

Problem Statement

leetcode problem link

Solution [TLE]

class Solution:
    def minCost(self, nums: List[int], cost: List[int]) -> int:
        res = float('inf')
        def normalize(x, arr):
            curr_cost = 0
            for i, num in enumerate(arr):
                delta = abs(num - x)
                curr_cost += (cost[i] * delta)
            return curr_cost

        for x in nums:
            res = min(res, normalize(x, nums[:]))
        return res

Editorial

Approach 1: Prefix Sum

class Solution:
    def minCost(self, nums: List[int], cost: List[int]) -> int:
        # Sort integers by values.
        num_and_cost = sorted([num, c] for num, c in zip(nums, cost))
        n = len(cost)

        # Get the prefix sum array of the costs.
        prefix_cost = [0] * n
        prefix_cost[0] = num_and_cost[0][1]
        for i in range(1, n):
            prefix_cost[i] = num_and_cost[i][1] + prefix_cost[i - 1]

        # Then we try every integer nums[i] and make every element equals nums[i],
        # Start with nums[0]
        total_cost = 0
        for i in range(1, n):
            total_cost += num_and_cost[i][1] * (num_and_cost[i][0] - num_and_cost[0][0])
        answer = total_cost

        # Then we try nums[1], nums[2] and so on. The cost difference is made by the change of
        # two parts: 1. prefix sum of costs. 2. suffix sum of costs.
        # During the iteration, record the minimum cost we have met.
        for i in range(1, n):
            gap = num_and_cost[i][0] - num_and_cost[i - 1][0]
            total_cost += prefix_cost[i - 1] * gap
            total_cost -= gap * (prefix_cost[n - 1] - prefix_cost[i - 1])
            answer = min(answer, total_cost)

        return answer

from typing import List

class Solution:
    def minCost(self, nums: List[int], cost: List[int]) -> int:
        # Pair and sort by value
        pairs = sorted(zip(nums, cost))  # [(value, weight)]
        n = len(pairs)

        # Prefix sums of weights (costs)
        prefix_w = [0] * n
        prefix_w[0] = pairs[0][1]
        for i in range(1, n):
            prefix_w[i] = prefix_w[i - 1] + pairs[i][1]

        # Cost to make everything equal to the smallest value (pairs[0][0])
        total = 0
        base = pairs[0][0]
        for i in range(1, n):
            val, w = pairs[i]
            total += w * (val - base)

        ans = total

        # Sweep target across sorted values: from pairs[i-1][0] to pairs[i][0]
        for i in range(1, n):
            gap = pairs[i][0] - pairs[i - 1][0]
            left_weight = prefix_w[i - 1]
            right_weight = prefix_w[n - 1] - prefix_w[i - 1]

            # Update total:
            # Left side gets farther by gap; right side gets closer by gap
            total += left_weight * gap
            total -= right_weight * gap

            if total < ans:
                ans = total

        return ans

Key Insights

• The objective f(T) = Σ w_i

  x_i - T

  (with w_i = cost[i], x_i = nums[i]) is convex in T.

• The minimizer is any weighted median of x_i under weights w_i.
• We can compute the minimal cost by:
  1. Finding a weighted median directly, or
  2. Scanning sorted x_i and updating the total cost in O(1) per step using prefix sums (prefix-sum sweep).

These approaches are equivalent in result; the sweep method is constructive and interview-friendly.

---

Prefix-Sum Sweep Intuition

1. Sort pairs by value: pairs = sorted(zip(nums, cost)) → [(x0,w0), (x1,w1), …], x0 ≤ x1 ≤ …
2. Precompute prefix weights: prefixw[i] = Σ{k=0..i} w_k.
3. Initialize total cost at target = x0:
   • total = Σ_{i=1..n-1} w_i * (x_i - x0).
4. Sweep target forward from x{i-1} to x_i (gap = x_i - x{i-1}):
   • Left side (indices < i) gets farther by gap → +gap _ sum_left_weights = +gap _ prefix_w[i-1].
   • Right side (indices ≥ i) gets closer by gap → -gap _ sum_right_weights = -gap _ (prefix_w[n-1] - prefix_w[i-1]).
   • Update total and track the minimum.

Why it works:

• Moving T by a small amount changes each term linearly (±w_i * gap).
• Summing over left/right partitions yields an O(1) update at each breakpoint x_i.
• Because f(T) is convex and piecewise-linear with breakpoints at x_i, checking all x_i finds the global minimum.

---

Complexity

• Sorting: O(n log n)
• Prefix sums: O(n)
• Sweep: O(n)
• Space: O(n)