当前位置：首页 > article >正文

【C++二分查找前缀和】2333. 最小差值平方和|2011

article 2024/12/2 7:16:07

本文涉及的基础知识点

C++算法：前缀和、前缀乘积、前缀异或的原理、源码及测试用例包括课程视频
C++二分查找

LeetCode2333. 最小差值平方和

给你两个下标从 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 <= 10⁵
0 <= nums1[i], nums2[i] <= 10⁵
0 <= k1, k2 <= 10⁹

二分查找+前缀和

nums3 = |nums1[i]-nums2[i]|,其和为sum。
显然nums1[i]，加一减一，和nums2[i]减一加一完全相同。本题 $\iff$ 将nums3[i]最多减1k1+k2次。
有序映射mCnt记录nums3[i]各值的数量。
有序映射mPreCnt[i]记录nums3中值小于等于i的数量。
有序映射mPreSum[i]记录nums3中值小于等于i的和。
mCnt和mPreCnt mPreSum的key都只包括nums3中的元素和0。
Check(x)
参数返回：[0,nums3的最大值]
it = mCnt.upper(x)
–it
假定修改成功后，其和为：sum1 = mPreSum[it.first] + (N - it.second)*x;
return sum1 - sum <= k1+k2。
二分求出最小x。
如果有剩余次数，如果x 不是0，将部分x变成x-1。
然后求平方和。

代码

核心代码

template<class INDEX_TYPE>
class CBinarySearch
{
public:
	CBinarySearch(INDEX_TYPE iMinIndex, INDEX_TYPE iMaxIndex):m_iMin(iMinIndex),m_iMax(iMaxIndex) {}
	template<class _Pr>
	INDEX_TYPE FindFrist( _Pr pr)
	{
		auto left = m_iMin - 1;
		auto rightInclue = m_iMax;
		while (rightInclue - left > 1)
		{
			const auto mid = left + (rightInclue - left) / 2;
			if (pr(mid))
			{
				rightInclue = mid;
			}
			else
			{
				left = mid;
			}
		}
		return rightInclue;
	}
	template<class _Pr>
	INDEX_TYPE FindEnd( _Pr pr)
	{
		int leftInclude = m_iMin;
		int right = m_iMax + 1;
		while (right - leftInclude > 1)
		{
			const auto mid = leftInclude + (right - leftInclude) / 2;
			if (pr(mid))
			{
				leftInclude = mid;
			}
			else
			{
				right = mid;
			}
		}
		return leftInclude;
	}
protected:
	const INDEX_TYPE m_iMin, m_iMax;
};

class Solution {
		public:
			long long minSumSquareDiff(vector<int>& nums1, vector<int>& nums2, int k1, int k2) {
				map<int, int> mCnt, mPreCnt;
				map<int,long long> mPreSum;
				const int N = nums1.size();
				mCnt[0] = 0;				
				for (int i = 0; i < N; i++) {
					mCnt[abs(nums1[i] - nums2[i])]++;
				}
				mPreCnt[0] = 0;
				mPreSum[0] = 0;				
				for (auto [v, cnt] : mCnt) {
					mPreCnt[v] = mPreCnt.rbegin()->second + cnt;
					mPreSum[v] = mPreSum.rbegin()->second + (long long)v * cnt;
				}
				const long long sum = mPreSum.rbegin()->second;
				auto Need = [&](int x) {
					auto it = mPreCnt.upper_bound(x);
					--it;
					const auto sum1 = mPreSum[it->first] + (long long)x * (N - it->second);
					return sum -sum1;
				};
				auto Check = [&](int x) {return Need(x) <= k1 + k2; };
				 CBinarySearch<int> bs(0, mCnt.rbegin()->first);
				 int x1 = bs.FindFrist(Check);
				 if (0 == x1) { return 0; }
				 long long ans=0;
				 int cnt1 = 0;
				 for (auto [v, cnt] : mCnt) {
					 if (v >= x1)break;
					 cnt1 += cnt;
					 ans += (long long)v * v * cnt;
				 }
				 const int iRemain = k1 + k2 - Need(x1);
				 const int xcnt = N - cnt1 - iRemain;
				 ans += (long long)(x1 - 1) * (x1 - 1) * iRemain;
				 ans += (long long)x1 * x1 * xcnt;
				 return ans;
			}
		};

单元测试

vector<int> nums1, nums2;
		int k1,k2;
		TEST_METHOD(TestMethod11)
		{
			nums1 = { 1,2,3,4 }, nums2 = { 2,10,20,19 }, k1 = 0, k2 = 0;
			auto res = Solution().minSumSquareDiff(nums1, nums2, k1, k2);
			AssertEx(579LL, res);
		}
		TEST_METHOD(TestMethod12)
		{
			nums1 = { 1,4,10,12 }, nums2 = { 5,8,6,9 }, k1 = 1, k2 = 1;
			auto res = Solution().minSumSquareDiff(nums1, nums2, k1, k2);
			AssertEx(43LL, res);
		}

扩展阅读

我想对大家说的话
工作中遇到的问题，可以按类别查阅鄙人的算法文章，请点击《算法与数据汇总》。
学习算法：按章节学习《喜缺全书算法册》，大量的题目和测试用例，打包下载。重视操作
有效学习：明确的目标及时的反馈拉伸区（难度合适）专注
闻缺陷则喜(喜缺)是一个美好的愿望，早发现问题，早修改问题，给老板节约钱。
子墨子言之：事无终始，无务多业。也就是我们常说的专业的人做专业的事。
如果程序是一条龙，那算法就是他的是睛
失败+反思=成功成功+反思=成功