Medium

题目描述

给你两个长度为 n 的正整数数组 nums1nums2

数组 nums1nums2 的平方差之和定义为所有 (nums1[i] - nums2[i])² 的和,其中 0 <= i < n

同时给你两个正整数 k1k2。你可以将 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 - 10)² + (3 - 20)² + (4 - 19)² = 579。

示例 2:

输入:nums1 = [1,4,10,12], nums2 = [5,8,6,9], k1 = 1, k2 = 1
输出:43
解释:一种获得最小平方差之和的方式是:
- 将 nums1[0] 增加一次。
- 将 nums2[2] 增加一次。
最小平方差之和为:(2 - 5)² + (4 - 8)² + (10 - 7)² + (12 - 9)² = 43。

约束条件:

  • n == nums1.length == nums2.length
  • 1 <= n <= 10⁵
  • 0 <= nums1[i], nums2[i] <= 10⁵
  • 0 <= k1, k2 <= 10⁹

解题思路

这道题的关键思路是贪心算法 + 二分搜索

首先理解问题本质:对 nums1[i] 加1等价于对 nums2[i] 减1,反之亦然。因此 k1k2 可以合并为总操作次数 k = k1 + k2

核心策略是贪心地减少最大的差值。每次操作都应该针对当前差值最大的位置,这样能最大化减少平方差之和。

解法思路:

  1. 计算初始差值:对每个位置计算 |nums1[i] - nums2[i]|
  2. 二分搜索最大差值:在操作 k 次后,剩余的最大差值是多少
  3. 贪心分配操作:优先处理差值大的位置,将其降到目标值

具体实现中,我们二分搜索最终的最大差值 maxDiff。对于每个可能的 maxDiff,计算需要多少次操作才能使所有差值都不超过 maxDiff。如果所需操作次数不超过 k,说明这个 maxDiff 是可达的。

找到最小的可达 maxDiff 后,我们模拟操作过程:先将所有大于 maxDiff 的差值降到 maxDiff,然后用剩余操作次数进一步优化。

代码实现

class Solution {
public:
    long long minSumSquareDiff(vector<int>& nums1, vector<int>& nums2, int k1, int k2) {
        int n = nums1.size();
        vector<long long> diffs(n);
        for (int i = 0; i < n; i++) {
            diffs[i] = abs(nums1[i] - nums2[i]);
        }
        
        long long k = (long long)k1 + k2;
        long long left = 0, right = *max_element(diffs.begin(), diffs.end());
        
        // 二分搜索最终的最大差值
        while (left < right) {
            long long mid = left + (right - left) / 2;
            long long need = 0;
            for (long long diff : diffs) {
                if (diff > mid) {
                    need += diff - mid;
                }
            }
            if (need <= k) {
                right = mid;
            } else {
                left = mid + 1;
            }
        }
        
        long long maxDiff = left;
        long long used = 0;
        
        // 将所有大于 maxDiff 的差值降到 maxDiff
        for (int i = 0; i < n; i++) {
            if (diffs[i] > maxDiff) {
                used += diffs[i] - maxDiff;
                diffs[i] = maxDiff;
            }
        }
        
        // 用剩余操作次数进一步优化
        long long remaining = k - used;
        if (maxDiff > 0) {
            long long count = 0;
            for (long long diff : diffs) {
                if (diff == maxDiff) count++;
            }
            
            long long reduce = min(remaining / count, maxDiff);
            remaining -= reduce * count;
            
            for (int i = 0; i < n; i++) {
                if (diffs[i] == maxDiff) {
                    diffs[i] -= reduce;
                    if (remaining > 0 && diffs[i] > 0) {
                        diffs[i]--;
                        remaining--;
                    }
                }
            }
        }
        
        long long result = 0;
        for (long long diff : diffs) {
            result += diff * diff;
        }
        return result;
    }
};
class Solution:
    def minSumSquareDiff(self, nums1: List[int], nums2: List[int], k1: int, k2: int) -> int:
        n = len(nums1)
        diffs = [abs(nums1[i] - nums2[i]) for i in range(n)]
        k = k1 + k2
        
        # 二分搜索最终的最大差值
        left, right = 0, max(diffs)
        while left < right:
            mid = (left + right) // 2
            need = sum(max(0, diff - mid) for diff in diffs)
            if need <= k:
                right = mid
            else:
                left = mid + 1
        
        max_diff = left
        used = 0
        
        # 将所有大于 max_diff 的差值降到 max_diff
        for i in range(n):
            if diffs[i] > max_diff:
                used += diffs[i] - max_diff
                diffs[i] = max_diff
        
        # 用剩余操作次数进一步优化
        remaining = k - used
        if max_diff > 0:
            count = sum(1 for diff in diffs if diff == max_diff)
            reduce = min(remaining // count, max_diff)
            remaining -= reduce * count
            
            for i in range(n):
                if diffs[i] == max_diff:
                    diffs[i] -= reduce
                    if remaining > 0 and diffs[i] > 0:
                        diffs[i] -= 1
                        remaining -= 1
        
        return sum(diff * diff for diff in diffs)
public class Solution {
    public long MinSumSquareDiff(int[] nums1, int[] nums2, int k1, int k2) {
        int n = nums1.Length;
        long[] diffs = new long[n];
        for (int i = 0; i < n; i++) {
            diffs[i] = Math.Abs(nums1[i] - nums2[i]);
        }
        
        long k = (long)k1 + k2;
        long left = 0, right = diffs.Max();
        
        // 二分搜索最终的最大差值
        while (left < right) {
            long mid = left + (right - left) / 2;
            long need = 0;
            foreach (long diff in diffs) {
                if (diff > mid) {
                    need += diff - mid;
                }
            }
            if (need <= k) {
                right = mid;
            } else {
                left = mid + 1;
            }
        }
        
        long maxDiff = left;
        long used = 0;
        
        // 将所有大于 maxDiff 的差值降到 maxDiff
        for (int i = 0; i < n; i++) {
            if (diffs[i] > maxDiff) {
                used += diffs[i] - maxDiff;
                diffs[i] = maxDiff;
            }
        }
        
        // 用剩余操作次数进一步优化
        long remaining = k - used;
        if (maxDiff > 0) {
            long count = diffs.Count(diff => diff == maxDiff);
            long reduce = Math.Min(remaining / count, maxDiff);
            remaining -= reduce * count;
            
            for (int i = 0; i < n; i++) {
                if (diffs[i] == maxDiff) {
                    diffs[i] -= reduce;
                    if (remaining > 0 && diffs[i] > 0) {
                        diffs[i]--;
                        remaining--;
                    }
                }
            }
        }
        
        long result = 0;
        foreach (long diff in diffs) {
            result += diff * diff;
        }
        return result;
    }
}
var minSumSquareDiff = function(nums1, nums2, k1, k2) {
    const n = nums1.length;
    const diffs = [];
    
    for (let i = 0; i < n; i++) {
        diffs.push(Math.abs(nums1[i] - nums2[i]));
    }
    
    diffs.sort((a, b) => b - a);
    
    let totalK = k1 + k2;
    
    for (let i = 0; i < n && totalK > 0; i++) {
        if (diffs[i] === 0) break;
        
        let nextVal = i + 1 < n ? diffs[i + 1] : 0;
        let reduction = diffs[i] - nextVal;
        let count = i + 1;
        let totalReduction = reduction * count;
        
        if (totalReduction <= totalK) {
            for (let j = 0; j <= i; j++) {
                diffs[j] = nextVal;
            }
            totalK -= totalReduction;
        } else {
            let canReduce = Math.floor(totalK / count);
            let remainder = totalK % count;
            
            for (let j = 0; j <= i; j++) {
                diffs[j] -= canReduce;
                if (j < remainder) {
                    diffs[j]--;
                }
            }
            totalK = 0;
        }
    }
    
    let result = 0;
    for (let diff of diffs) {
        result += diff * diff;
    }
    
    return result;
};

复杂度分析

复杂度类型复杂度说明
时间复杂度O(n log max_diff)二分搜索需要 log max_diff 次,每次需要 O(n) 时间遍历数组
空间复杂度O(n)存储差值数组需要 O(n) 空间

相关题目