Problem of The Day: Minimum Operations to Make All Array Elements Equal

4 min read 831 words

Problem Statement

leetcode problem link

Solution [TLE]

class Solution:
    def minOperations(self, nums: List[int], queries: List[int]) -> List[int]:
        res = []
        for target in queries:
            curr = 0
            for num in nums:
                curr += abs(target - num)
            res.append(curr)
        return res

Solution [Accepted]

Algorithm for bisect_right manually Definition: bisect_right(a, x) returns the insertion index to place x after any existing entries equal to x (i.e., first index > x). It’s also the count of elements ≤ x.

# Python
def bisect_right(a, x):
    lo, hi = 0, len(a)
    while lo < hi:
        mid = (lo + hi) // 2
        if a[mid] <= x:
            lo = mid + 1   # move right if mid value <= x
        else:
            hi = mid       # keep mid in the range
    return lo

from typing import List
import bisect

class Solution:
    def minOperations(self, nums: List[int], queries: List[int]) -> List[int]:
        nums.sort()
        n = len(nums)
        ps = [0] * (n + 1)
        for i in range(n):
            ps[i + 1] = ps[i] + nums[i]

        res = []
        for x in queries:
            k = bisect.bisect_right(nums, x)  # count of elements <= x
            left_cost = x * k - ps[k]
            right_cost = (ps[n] - ps[k]) - x * (n - k)
            res.append(left_cost + right_cost)
        return res

Why Sorting + Prefix Sums Works

Let nums be sorted. Let ps[i] be the prefix sum of the first i elements (1-based), with ps[0] = 0.

For a query value x:

• Let k = upper_bound(nums, x) (i.e., number of elements <= x).
• Left side: indices [0 .. k-1] (k elements)
  • Cost to raise them to x: x*k - (sum of left) = x*k - ps[k].
• Right side: indices [k .. n-1] (n - k elements)
  • Cost to lower them to x: (sum of right) - x*(n-k) = (ps[n] - ps[k]) - x*(n-k).

Total cost = left_cost + right_cost.

Key properties:

• Binary search gives k in O(log n).
• Prefix sums give segment sums in O(1).
• Preprocessing is O(n log n) for sort + O(n) for prefix, queries are O(log n) each.

---

Complexity

• Preprocessing:
  • Sort: O(n log n)
  • Prefix sums: O(n)
• Per query: O(log n) (binary search)
• Space: O(n) for prefix sums

This improves over the naive O(n) per query to O(log n) per query, which is significant when there are many queries.

---

Worked Example

nums = [3, 1, 6, 8], queries = [1, 5]

1. Sort: nums = [1, 3, 6, 8]
   Prefix sums ps = [0, 1, 4, 10, 18]

Query x = 5:

• k = upper_bound([1,3,6,8], 5) = 2 (elements <= 5 are [1,3])
• Left cost = 5*2 - ps[2] = 10 - 4 = 6 (raise 1→5 cost 4, 3→5 cost 2)
• Right cost = (ps[4] - ps[2]) - 5*(4-2) = (18 - 4) - 10 = 4 (lower 6→5 cost 1, 8→5 cost 3)
• Total = 6 + 4 = 10

Query x = 1:

• k = upper_bound(…, 1) = 1
• Left cost = 1*1 - ps[1] = 1 - 1 = 0
• Right cost = (ps[4] - ps[1]) - 1*(4-1) = (18 - 1) - 3 = 14
• Total = 14