C++二分查找算法：有序矩阵中的第 k 个最小数组和

本文涉及的基础知识点

二分查找算法合集

本题的简化

C++二分查找算法：查找和最小的 K 对数字十分接近m恒等于2

题目

给你一个 m * n 的矩阵 mat，以及一个整数 k ，矩阵中的每一行都以非递减的顺序排列。
你可以从每一行中选出 1 个元素形成一个数组。返回所有可能数组中的第 k 个最小数组和。
示例 1：
输入：mat = [[1,3,11],[2,4,6]], k = 5
输出：7
解释：从每一行中选出一个元素，前 k 个和最小的数组分别是：
[1,2], [1,4], [3,2], [3,4], [1,6]。其中第 5 个的和是 7 。
示例 2：
输入：mat = [[1,3,11],[2,4,6]], k = 9
输出：17
示例 3：
输入：mat = [[1,10,10],[1,4,5],[2,3,6]], k = 7
输出：9
解释：从每一行中选出一个元素，前 k 个和最小的数组分别是：
[1,1,2], [1,1,3], [1,4,2], [1,4,3], [1,1,6], [1,5,2], [1,5,3]。其中第 7 个的和是 9 。
示例 4：
输入：mat = [[1,1,10],[2,2,9]], k = 7
输出：12
参数范围：
m == mat.length
n == mat.length[i]
1 <= m, n <= 40
1 <= k <= min(200, n ^ m)
1 <= mat[i][j] <= 5000
mat[i] 是一个非递减数组

分析

时间复杂度

O(mlog(500040)n+mkn)。GetLessKSum被调用m次，GetLessEqualSumNum共被调用mlog(500040)次。每次调用GetLessEqualSumNum，for循环共执行m次。
vRet.emplace_back极端情况下，可能被执行kn次。

主要函数介绍


GetLessKSum	两行升序数据的最小k个和
GetLessEqualSumNum	两行升序数据和小于等于iSum的组合数量

注意：nums[i]为正数，所以如果pre的数量大于k，只需要保留前k小，其它的被淘汰了。

二分

寻找第一个符合条件的iSum，条件如下：
和小于等于iSum的组合数量大于等于k。

代码

核心代码

class Solution {
public:
	int kthSmallest(vector<vector<int>>& mat, int k) {
		m_c = mat.front().size();
		m_iK = k;
		vector<int> pre = mat[0];
		for (int r = 1; r < mat.size(); r++)
		{
			pre = GetLessKSum(pre, mat[r]);
		}
		return pre.back();
	}
	vector<int> GetLessKSum(const vector<int>& pre, const vector<int>& cur)
	{
		int left = 0, right = 5000 * 40;
		while (right - left > 1)
		{
			const auto mid = left + (right - left) / 2;
			if (GetLessEqualSumNum(pre, cur, mid)>= m_iK)
			{
				right = mid;
			}
			else
			{
				left = mid;
			}
		}
		vector<int> vRet;
		for (const auto& pr : pre)
		{
			for (const auto& cu : cur)
			{
				if (pr + cu <= right)
				{
					vRet.emplace_back(pr + cu);
				}
				else
				{
					break;
				}
			}
		}
		sort(vRet.begin(), vRet.end());
		if (vRet.size() > m_iK)
		{
			vRet.erase(vRet.begin() + m_iK, vRet.end());
		}
		return vRet;
	}
	int GetLessEqualSumNum(const vector<int>& pre, const vector<int>& cur,int iSum)
	{
		int iNum = 0;
		for (const auto& pr : pre)
		{
			iNum += std::upper_bound(cur.begin(), cur.end(), iSum - pr)- cur.begin();
		}
		return iNum;
	}
	int m_iK;
	int m_c;
};

测试用例

template
void Assert(const T& t1, const T& t2)
{
assert(t1 == t2);
}

template
void Assert(const vector& v1, const vector& v2)
{
if (v1.size() != v2.size())
{
assert(false);
return;
}
for (int i = 0; i < v1.size(); i++)
{
Assert(v1[i], v2[i]);
}
}

int main()
{
vector<vector> mat;
int k;
int res;
{
Solution slu;
mat = { {1,3,11},{2,4,6} };
k = 5;
res = slu.kthSmallest(mat, k);
Assert(7, res);
}
{
Solution slu;
mat = { {1,3,11},{2,4,6} };
k = 9;
res = slu.kthSmallest(mat, k);
Assert(17, res);
}
{
Solution slu;
mat = { {1,10,10},{1,4,5},{2,3,6} };
k = 7;
res = slu.kthSmallest(mat, k);
Assert(9, res);
}
{
Solution slu;
mat = { {1,1,10},{2,2,9} };
k = 7;
res = slu.kthSmallest(mat, k);
Assert(12, res);
}

//CConsole::Out(res);

}

优化增加结果

vector<int> vRet;
	for (const auto& pr : pre)
	{
		for (const auto& cu : cur)
		{
			if (pr + cu < right)
			{
				vRet.emplace_back(pr + cu);
			}
			else
			{
				break;
			}
		}
	}
	while (vRet.size() < m_iK)
	{
		vRet.emplace_back(right);
	}

和小于right的数量<=k，如果不足够，则补right。时间复杂度由O(nk)降低到O(k+n)。

直接使用封装类

namespace NBinarySearch
{
	template<class INDEX_TYPE,class _Pr>
	INDEX_TYPE FindFrist(INDEX_TYPE left, INDEX_TYPE right, _Pr pr)
	{
		while (right - left > 1)
		{
			const auto mid = left + (right - left) / 2;
			if (pr(mid))
			{
				right = mid;
			}
			else
			{
				left = mid;
			}
		}
		return right;
	}
}


class Solution {
public:
	int kthSmallest(vector<vector<int>>& mat, int k) {
		m_c = mat.front().size();
		m_iK = k;
		vector<int> pre = mat[0];
		for (int r = 1; r < mat.size(); r++)
		{
			pre = GetLessKSum(pre, mat[r]);
		}
		return pre.back();
	}
	vector<int> GetLessKSum(const vector<int>& pre, const vector<int>& cur)
	{
		auto GetLessEqualSumNum = [&pre, &cur, this](const int iSum)-> bool
		{
			int iNum = 0;
			for (const auto& pr : pre)
			{
				iNum += std::upper_bound(cur.begin(), cur.end(), iSum - pr) - cur.begin();
			}
			return iNum >= m_iK;
		};
		const int right = NBinarySearch::FindFrist(0, 5000 * 40, GetLessEqualSumNum);		
		vector<int> vRet;
		for (const auto& pr : pre)
		{
			for (const auto& cu : cur)
			{
				if (pr + cu < right)
				{
					vRet.emplace_back(pr + cu);
				}
				else
				{
					break;
				}
			}
		}
		while (vRet.size() < m_iK)
		{
			vRet.emplace_back(right);
		}
		sort(vRet.begin(), vRet.end());
		if (vRet.size() > m_iK)
		{
			vRet.erase(vRet.begin() + m_iK, vRet.end());
		}
		return vRet;
	}
	int GetLessEqualSumNum(const vector<int>& pre, const vector<int>& cur,int iSum)
	{
		int iNum = 0;
		for (const auto& pr : pre)
		{
			iNum += std::upper_bound(cur.begin(), cur.end(), iSum - pr)- cur.begin();
		}
		return iNum;
	}
	int m_iK;
	int m_c;
};

2023年3月暴力版

直接保留前k个。时间复杂度:O(mknlogk)
class Solution {
public:
int kthSmallest(vector<vector>& mat, int k) {
m_r = mat.size();
m_c = mat[0].size();
std::priority_queue pre;
pre.push(0);
for (int r = 0; r < mat.size(); r++)
{
std::priority_queue dp;
while (pre.size())
{
int t = pre.top();
pre.pop();
for (int c = 0; c < m_c; c++)
{
dp.push(mat[r][c] + t);
if (dp.size() > k)
{
dp.pop();
}
}
}
pre.swap(dp);
}
return pre.top();
}
int m_r, m_c;
};

扩展阅读

视频课程

有效学习：明确的目标及时的反馈拉伸区（难度合适），可以先学简单的课程，请移步CSDN学院，听白银讲师（也就是鄙人）的讲解。
https://edu.csdn.net/course/detail/38771

如何你想快

速形成战斗了，为老板分忧，请学习C#入职培训、C++入职培训等课程
https://edu.csdn.net/lecturer/6176

我想对大家说的话
闻缺陷则喜是一个美好的愿望，早发现问题，早修改问题，给老板节约钱。
墨子曰：事无终始，无务多业。也就是我们常说的专业的人做专业的事。
如果程序是一条龙，那算法就是他的是睛