跳过正文
  1. leetcode 题解/

2602_使数组元素全部相等的最少操作次数

·307 字·2 分钟
目录

题目
#

给你一个正整数数组 nums 。

同时给你一个长度为 m 的整数数组 queries 。第 i 个查询中,你需要将 nums 中所有元素变成 queries[i] 。你可以执行以下操作 任意 次:

  • 将数组里一个元素 增大 或者 减小 1 。

请你返回一个长度为 m 的数组 **answer ,其中 **answer[i]是将 nums 中所有元素变成 queries[i] 的 最少 操作次数。

注意,每次查询后,数组变回最开始的值。

示例 1:

输入:nums = [3,1,6,8], queries = [1,5]
输出:[14,10]
解释:第一个查询,我们可以执行以下操作:
- 将 nums[0] 减小 2 次,nums = [1,1,6,8] 。
- 将 nums[2] 减小 5 次,nums = [1,1,1,8] 。
- 将 nums[3] 减小 7 次,nums = [1,1,1,1] 。
第一个查询的总操作次数为 2 + 5 + 7 = 14 。
第二个查询,我们可以执行以下操作:
- 将 nums[0] 增大 2 次,nums = [5,1,6,8] 。
- 将 nums[1] 增大 4 次,nums = [5,5,6,8] 。
- 将 nums[2] 减小 1 次,nums = [5,5,5,8] 。
- 将 nums[3] 减小 3 次,nums = [5,5,5,5] 。
第二个查询的总操作次数为 2 + 4 + 1 + 3 = 10 。

示例 2:

输入:nums = [2,9,6,3], queries = [10]
输出:[20]
解释:我们可以将数组中所有元素都增大到 10 ,总操作次数为 8 + 1 + 4 + 7 = 20 。

提示:

  • n == nums.length
  • m == queries.length
  • 1 <= n, m <= 105
  • 1 <= nums[i], queries[i] <= 109

解法
#

我们对 nums 升序排列,然后计算前缀和。

对于每次 query, value = queries[i] 我们通过二分法找到第一个 >= value 的数的下标,那么该位置左边的数均 < value, 该位置及右边的数均 >= value.

通过前缀和,我们知道左边要补多少,右边要减多少。

代码
#

class Solution:
    def binSearch(self, nums, target, left, right):
        if right < left:
            return left
        mid = (left + right) // 2
        if nums[mid] < target:
            return self.binSearch(nums, target, mid + 1, right)
        else:
            return self.binSearch(nums, target, left, mid - 1)

    def minOperations(self, nums: List[int], queries: List[int]) -> List[int]:
        nums.sort()
        n = len(nums)
        prefix = nums[:]
        for i in range(1, n):
            prefix[i] += prefix[i-1]
        ans = []
        for query in queries:
            idx = self.binSearch(nums, query, 0, n-1)
            ops = 0
            if idx == 0 or idx == n-1:
                ops = abs(prefix[-1] - query * n)
            else:
                ops = (query * idx - prefix[idx-1]) + (prefix[-1] - prefix[idx-1] - query * (n-idx))
            ans.append(ops)
        return ans