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

Raphael Liu Lv10
  2023-09-13 00:47:26 2023-09-13 00:47:26 Created   2023-09-13 17:01:06 2023-09-13 17:01:06 Updated    

给你一个正整数数组 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

本题视频讲解

见【周赛 338】 。

前置知识：前缀和

对于数组 nums，定义它的前缀和 s}[0]=0，s}[i+1] = \sum\limits_{j=0}^{i}\textit{nums}[j]。

根据这个定义，有 s[i+1]=s[i]+\textit{nums}[i]。

例如 nums}=[1,2,-1,2]，对应的前缀和数组为 s=[0,1,3,2,4]。

通过前缀和，我们可以把子数组的元素和转换成两个前缀和的差，即

\sum_{j=\textit{left} }^{\textit{right} }\textit{nums}[j] = \sum\limits_{j=0}^{\textit{right} }\textit{nums}[j] - \sum\limits_{j=0}^{\textit{left}-1}\textit{nums}[j] = \textit{s}[\textit{right}+1] - \textit{s}[\textit{left}]

例如 nums 的子数组 [2,-1,2] 的和就可以用 s[4]-s[1]=4-1=3 算出来。

注：为方便计算，常用左闭右开区间 [\textit{left},\textit{right}) 来表示从 nums}[\textit{left}] 到 nums}[\textit{right}-1] 的子数组，此时子数组的和为 s}[\textit{right}] - \textit{s}[\textit{left}]，子数组的长度为 right}-\textit{left。

注 2：s[0]=0 表示一个空数组的元素和。为什么要额外定义它？想一想，如果要计算的子数组恰好是一个前缀（从 nums}[0] 开始），你要用 s[\textit{right}] 减去谁呢？通过定义 s[0]=0，任意子数组（包括前缀）都可以表示为两个前缀和的差。

前置知识：二分查找

见【基础算法精讲 04】 。

思路

t3.png

[sol1-Python3]
1
2
3
4
5
6
7
8
9
10
11
12
class Solution:
    def minOperations(self, nums: List[int], queries: List[int]) -> List[int]:
        n = len(nums)
        nums.sort()
        s = list(accumulate(nums, initial=0))  # 前缀和
        ans = []
        for q in queries:
            j = bisect_left(nums, q)
            left = q * j - s[j]  # 蓝色面积
            right = s[n] - s[j] - q * (n - j)  # 绿色面积
            ans.append(left + right)
        return ans
[sol1-Java]
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
28
29
30
31
32
33
34
class Solution {
    public List<Long> minOperations(int[] nums, int[] queries) {
        Arrays.sort(nums);
        int n = nums.length;
        var sum = new long[n + 1]; // 前缀和
        for (int i = 0; i < n; ++i)
            sum[i + 1] = sum[i] + nums[i];

        var ans = new ArrayList<Long>(queries.length);
        for (int q : queries) {
            int j = lowerBound(nums, q);
            long left = (long) q * j - sum[j]; // 蓝色面积
            long right = sum[n] - sum[j] - (long) q * (n - j); // 绿色面积
            ans.add(left + right);
        }
        return ans;
    }

    // 见 https://www.bilibili.com/video/BV1AP41137w7/
    private int lowerBound(int[] nums, int target) {
        int left = -1, right = nums.length; // 开区间 (left, right)
        while (left + 1 < right) { // 区间不为空
            // 循环不变量：
            // nums[left] < target
            // nums[right] >= target
            int mid = left + (right - left) / 2;
            if (nums[mid] < target)
                left = mid; // 范围缩小到 (mid, right)
            else
                right = mid; // 范围缩小到 (left, mid)
        }
        return right;
    }
}
[sol1-C++]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
class Solution {
public:
    vector<long long> minOperations(vector<int> &nums, vector<int> &queries) {
        sort(nums.begin(), nums.end());
        int n = nums.size();
        long long sum[n + 1]; // 前缀和
        sum[0] = 0;
        for (int i = 0; i < n; ++i)
            sum[i + 1] = sum[i] + nums[i];

        int m = queries.size();
        vector<long long> ans(m);
        for (int i = 0; i < m; ++i) {
            int q = queries[i];
            long long j = lower_bound(nums.begin(), nums.end(), q) - nums.begin();
            long long left = q * j - sum[j]; // 蓝色面积
            long long right = sum[n] - sum[j] - q * (n - j); // 绿色面积
            ans[i] = left + right;
        }
        return ans;
    }
};
[sol1-Go]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
func minOperations(nums, queries []int) []int64 {
	n := len(nums)
	sort.Ints(nums)
	sum := make([]int, n+1) // 前缀和
	for i, x := range nums {
		sum[i+1] = sum[i] + x
	}
	ans := make([]int64, len(queries))
	for i, q := range queries {
		j := sort.SearchInts(nums, q)
		left := q*j - sum[j] // 蓝色面积
		right := sum[n] - sum[j] - q*(n-j) // 绿色面积
		ans[i] = int64(left + right)
	}
	return ans
}

复杂度分析

  • 时间复杂度：O((n+q)\log n)，其中 n 为 nums 的长度，q 为 queries 的长度。
  • 空间复杂度：O(n)。返回值不计入。
 Comments
On this page
2602-使数组元素全部相等的最少操作次数
  1. 本题视频讲解
  2. 前置知识：前缀和
  3. 前置知识：二分查找
  4. 思路
  5. 复杂度分析