Program to find minimum total cost for equalizing list elements in Python

Suppose we have two lists of numbers called nums and costs. We can increase or decrease nums[i] for cost costs[i]. We want to make all elements in nums equal with minimum total cost.

So, if the input is like nums = [3, 2, 4] and costs = [1, 10, 2], then the output will be 5. We can decrease 3 to 2 for cost 1, then decrease 4 to 2 (twice) for cost 2 each, totaling 1 + 2 + 2 = 5.

Approach

This problem uses ternary search to find the optimal target value. The cost function is unimodal (has a single minimum), making ternary search ideal ?

Algorithm Steps

• Define a helper function to calculate total cost for a given target

• Use ternary search between 0 and maximum value in nums

• Compare costs at mid and mid+1 to narrow the search range

• Return the minimum cost found

Implementation

class Solution:
    def solve(self, nums, costs):
        def helper(target):
            total = 0
            for i, n in enumerate(nums):
                if target != n:
                    total += abs(n - target) * costs[i]
            return total
        
        low, high = 0, max(nums)
        while low < high:
            mid = (low + high) // 2
            if helper(mid) < helper(mid + 1):
                high = mid
            else:
                low = mid + 1
        return helper(low)

# Test the solution
ob = Solution()
nums = [3, 2, 4]
costs = [1, 10, 2]
print(ob.solve(nums, costs))

5

How It Works

The algorithm works by finding the optimal target value that minimizes total cost ?

• Helper function: Calculates total cost to make all elements equal to target

• Ternary search: Finds the target value with minimum cost efficiently

• Cost calculation: For each element, cost = |current_value - target| × cost_per_operation

Example Walkthrough

For nums = [3, 2, 4] and costs = [1, 10, 2] with target = 2 ?

• Element 3: |3-2| × 1 = 1

• Element 2: |2-2| × 10 = 0

• Element 4: |4-2| × 2 = 4

• Total cost: 1 + 0 + 4 = 5

Conclusion

This solution uses ternary search to efficiently find the optimal target value in O(n log max_value) time complexity. The key insight is that the cost function is unimodal, having exactly one minimum point.