[LeetCode]Find K Pairs with Smallest Sums

如果我们将两个数组展开，可以得到一个2D的matrix。例如[1,2,3]和[1,1,2]。我们可以得到下图：

我们可以看出来，这个matrix从左向右，自上而下是递增的。假设有m行，n列，就可以看做m个长度为n的sorted list，那么这个问题也就转变为merge m sorted list and get first k element，那么很明显我们就可以用priority queue来解。我们先把[0, 0]加入优先队列，每次从队列中取出最小的，如果其下边和右边还有元素就将其推入优先队列，直到得到k个元素。循环只有k次，每次我们取出一个最多推入两个，优先队列初始长度为1，所以队列长度<= k。时间按复杂度O(k)，空间复杂度O(k)，代码如下：

	class Solution {
	public:
	vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
	vector<pair<int, int>> res;
	int len1 = nums1.size(), len2 = nums2.size();
	auto comp = [&nums1, &nums2](const pair<int, int>& p1, const pair<int, int>& p2){return nums1[p1.first] + nums2[p1.second] > nums1[p2.first] + nums2[p2.second];};
	priority_queue<pair<int, int>, vector<pair<int, int>>, decltype(comp)> pq(comp);
	if(len1 && len2)pq.push({0, 0});
	while(--k >= 0 && pq.size())
	{
	auto p = pq.top();
	pq.pop();
	res.push_back({nums1[p.first], nums2[p.second]});
	if(!p.first && p.second < len2 - 1)
	pq.push({p.first, p.second + 1});
	if(p.first < len1 - 1)
	pq.push({p.first + 1, p.second});
	}
	return res;
	}
	};

view raw Find K Pairs with Smallest Sums.cpp hosted with ❤ by GitHub

另外一点，如果这道题不是求前k个，而是求第k个，我们还可以用值域的二分法来做。例如，Kth Smallest Element in Sorted Array, Find Kth Smallest Pair Distance