LC 6118. 最小差值平方和
https://leetcode.cn/contest/biweekly-contest-82/problems/minimum-sum-of-squared-difference/
这题可以抽象成为下面这个问题
- 假设我们有个数组A,里面每个数值都是正数
- 我们有K次机会去对任意元素进行-1操作
- 使得 `sum((x * x for x in A))` 结果最小
其实就是怎么快速地平摊K到这些元素上面,然后让每个元素都尽可能低接近。一个方法是使用heap, 每次选出最大的元素减一,减少最大和次大之间的差值。不过如果这个K很大,那么就会耗时比较长。
解决方法其实和 这题 非常类似,思路就是找到一个参考点,以这个参考点进行计算。这个问题的参考点就是:最多能调整多少个元素。
下面是具体实现代码,里面有比较详细的注释。(这题我写了/调试了很长时间)
class Solution:
def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:
n = len(nums1)
diff = []
for i in range(n):
diff.append(abs(nums1[i] - nums2[i]))
# 如果可以全部调整完成
if sum(diff) <= (k1 + k2):
return 0
diff.sort(reverse=True)
tt = 0
index = 0
k = k1 + k2
# 如果前面index个元素全部下降到diff[index]的话,那么会超过k
# 说明我们最多调整前面index个元素
while index < n:
if (tt - (diff[index] * index)) > k:
break
tt += diff[index]
index += 1
# 前面index个元素调整的话,至少可以调整到
# base = floor((tt - k) / index) = (tt - k + index - 1) // index
base = (tt - k + index - 1) // index
for i in range(index):
delta = diff[i] - base
diff[i] -= delta
k -= delta
for i in range(index):
if k > 0:
diff[i] -= 1
k -= 1
ans = sum((x * x for x in diff))
return ans