LC 6118. 最小差值平方和

https://leetcode.cn/contest/biweekly-contest-82/problems/minimum-sum-of-squared-difference/

这题可以抽象成为下面这个问题

其实就是怎么快速地平摊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