2333-最小差值平方和

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

给你两个下标从 0 开始的整数数组 nums1 和 nums2 ，长度为 n 。

数组 nums1 和 nums2 的 差值平方和 定义为所有满足 0 <= i < n 的 (nums1[i] - nums2[i])2 之和。

同时给你两个正整数 k1 和 k2 。你可以将 nums1 中的任意元素 +1 或者 -1 至多 k1 次。类似的，你可以将
nums2 中的任意元素 +1 或者 -1 至多 k2 次。

请你返回修改数组 _ _nums1 _ _ 至多 _ _k1 次且修改数组 _ _nums2 至多 k2 _ _ 次后的最小
差值平方和 。

注意： 你可以将数组中的元素变成 负 整数。

示例 1：

**输入：** nums1 = [1,2,3,4], nums2 = [2,10,20,19], k1 = 0, k2 = 0
**输出：** 579
**解释：** nums1 和 nums2 中的元素不能修改，因为 k1 = 0 和 k2 = 0 。
差值平方和为：(1 - 2)2 + (2 - 10)2 + (3 - 20)2 + (4 - 19)2 = 579 。

示例 2：

**输入：** nums1 = [1,4,10,12], nums2 = [5,8,6,9], k1 = 1, k2 = 1
**输出：** 43
**解释：** 一种得到最小差值平方和的方式为：
- 将 nums1[0] 增加一次。
- 将 nums2[2] 增加一次。
最小差值平方和为：
(2 - 5)2 + (4 - 8)2 + (10 - 7)2 + (12 - 9)2 = 43 。
注意，也有其他方式可以得到最小差值平方和，但没有得到比 43 更小答案的方案。

提示：

  • n == nums1.length == nums2.length
  • 1 <= n <= 105
  • 0 <= nums1[i], nums2[i] <= 105
  • 0 <= k1, k2 <= 109

本题 视频讲解 已出炉，欢迎点赞三连~

根据题意，在 nums}_1[i] 上 +1，等价于在 nums}_2[i] 上 -1，反之亦然。

定义 a[i]=|\textit{nums}_1[i]-\textit{nums}_2[i]|，k=k_1+k_2，则原问题可以转换成：

对数组 a 执行至多 k 次 -1 操作，能得到的 \sum a[i]^2 的最小值。

对于两个数，先把大的 -1 会更优。我们可以将 a 从大往小排序，然后从左到右遍历 a，同时更新剩余操作次数 k。

当遍历至 a[i] 时，a[0] 到 a[i-1] 均已减小至 a[i]，我们需要判断 k 次操作能否让 a[0] 到 a[i] 全部减小至 a[i+1]，即比较 k 与所需次数 c = (i + 1) (a[i] - a[i+1]) 的大小：

  • 如果 c<k，那么从 a[0] 到 a[i] 均可以减小至 a[i+1]，更新 k=k-c。
  • 如果 c\ge k，那么从 a[0] 到 a[i] 中：
    • 有 k \bmod (i+1) 个元素可以额外减小 \left\lfloor\dfrac{k}{i+1}\right\rfloor+1；
    • 有 i+1-k \bmod (i+1) 个元素可以额外减小 \left\lfloor\dfrac{k}{i+1}\right\rfloor。
      后续无法继续减小，应退出循环。

代码实现时，可以在 a 末尾加一个哨兵 0，减少边界判断。

复杂度分析

  • 时间复杂度：O(n\log n)。瓶颈在排序上。
  • 空间复杂度：O(1)。忽略排序的栈开销和哨兵的开销。
[sol1-Python3]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
class Solution:
    def minSumSquareDiff(self, a: List[int], nums2: List[int], k1: int, k2: int) -> int:
        ans, k = 0, k1 + k2
        for i in range(len(a)):
            a[i] = abs(a[i] - nums2[i])
            ans += a[i] * a[i]
        if sum(a) <= k:
            return 0  # 所有 a[i] 均可为 0
        a.sort(reverse=True)
        a.append(0)  # 哨兵
        for i, v in enumerate(a):
            ans -= v * v
            j = i + 1
            c = j * (v - a[j])
            if c < k:
                k -= c
                continue
            v -= k // j
            return ans + k % j * (v - 1) * (v - 1) + (j - k % j) * v * v
[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
class Solution {
    public long minSumSquareDiff(int[] a, int[] nums2, int k1, int k2) {
        int n = a.length, k = k1 + k2;
        long ans = 0L, sum = 0L;
        for (var i = 0; i < n; ++i) {
            a[i] = Math.abs(a[i] - nums2[i]);
            sum += a[i];
            ans += (long) a[i] * a[i];
        }
        if (sum <= k) return 0; // 所有 a[i] 均可为 0
        Arrays.sort(a);
        for (var i = n - 1; ; --i) {
            var m = n - i;
            long v = a[i], c = m * (v - (i > 0 ? a[i - 1] : 0));
            ans -= v * v;
            if (c < k) {
                k -= c;
                continue;
            }
            v -= k / m;
            return ans + k % m * (v - 1) * (v - 1) + (m - k % m) * v * v;
        }
    }
}
[sol1-C++]
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:
    long long minSumSquareDiff(vector<int> &a, vector<int> &nums2, int k1, int k2) {
        int n = a.size(), k = k1 + k2;
        long ans = 0L, sum = 0L;
        for (int i = 0; i < n; ++i) {
            a[i] = abs(a[i] - nums2[i]);
            sum += a[i];
            ans += (long) a[i] * a[i];
        }
        if (sum <= k) return 0; // 所有 a[i] 均可为 0
        sort(a.begin(), a.end(), greater<int>());
        a.push_back(0); // 哨兵
        for (int i = 0;; ++i) {
            long j = i + 1, v = a[i], c = j * (v - a[j]);
            ans -= v * v;
            if (c < k) {
                k -= c;
                continue;
            }
            v -= k / j;
            return ans + k % j * (v - 1) * (v - 1) + (j - k % j) * v * v;
        }
    }
};
[sol1-Go]
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
func minSumSquareDiff(a, nums2 []int, k1, k2 int) int64 {
	ans, sum := 0, 0
	for i, v := range a {
		a[i] = abs(v - nums2[i])
		sum += a[i]
		ans += a[i] * a[i]
	}
	k := k1 + k2
	if sum <= k {
		return 0 // 所有 a[i] 均可为 0
	}
	sort.Sort(sort.Reverse(sort.IntSlice(a)))
	a = append(a, 0) // 哨兵
	for i, v := range a {
		i++
		ans -= v * v
		if c := i * (v - a[i]); c < k {
			k -= c
			continue
		}
		v -= k / i
		ans += k%i*(v-1)*(v-1) + (i-k%i)*v*v
		break
	}
	return int64(ans)
}

func abs(x int) int { if x < 0 { return -x }; return x }
 Comments
On this page
2333-最小差值平方和