Skip to content

2602. Minimum Operations to Make All Array Elements Equal 👍

  • Time: $O(\texttt{sort})$
  • Space: $O(n)$
 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:
  vector<long long> minOperations(vector<int>& nums, vector<int>& queries) {
    const int n = nums.size();
    vector<long long> ans;
    vector<long long> prefix{0};

    ranges::sort(nums);

    for (int i = 0; i < n; ++i)
      prefix.push_back(prefix.back() + nums[i]);

    for (const long query : queries) {
      const int i = ranges::upper_bound(nums, query) - nums.begin();
      // Since nums[0..i) <= query, nums[i..n) > `query`, we should
      // - increase each num in nums[0..i) to `query` and
      // - decrease each num in nums[i..n) to `query`.
      const long inc = query * i - prefix[i];
      const long dec = prefix[n] - prefix[i] - query * (n - i);
      ans.push_back(inc + dec);
    }

    return ans;
  }
};

 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
class Solution {
  public List<Long> minOperations(int[] nums, int[] queries) {
    final int n = nums.length;
    List<Long> ans = new ArrayList<>();
    long[] prefix = new long[n + 1];

    Arrays.sort(nums);

    for (int i = 0; i < n; ++i)
      prefix[i + 1] = prefix[i] + nums[i];

    for (final int query : queries) {
      int i = Arrays.binarySearch(nums, query);
      if (i < 0)
        i = -i - 1;

      // Since nums[0..i) <= query, nums[i..n) > `query`, we should
      // - increase each num in nums[0..i) to `query` and
      // - decrease each num in nums[i..n) to `query`.
      final long inc = (long) query * i - prefix[i];
      final long dec = prefix[n] - prefix[i] - (long) query * (n - i);
      ans.add(inc + dec);
    }

    return ans;
  }
}

1
2
3
4
5
6
7
8
9
class Solution:
  def minOperations(self, nums: list[int], queries: list[int]) -> list[int]:
    n = len(nums)
    nums.sort()
    prefix = list(itertools.accumulate(nums, initial=0))
    splits = [(query, bisect.bisect_right(nums, query)) for query in queries]
    return [(query * i - prefix[i]) +
            (prefix[-1] - prefix[i] - query * (n - i))
            for query, i in splits]