求两个有序数组两两相加的值最小的K个数

2024年6月11日 126次阅读

题目描述：

有两个大小分别是lenA和lenB的数组A,B，它们元素的按非递减有序排列，找出这样的k个最小的(ai + bj) ,其中 0<= i < lenA,0<= j < lenB，要求算法的时间复杂度和空间复杂度尽量低。

例如对于：

A = 1,2,3,4

B = 2,3,4,5

ai+bj的所有组合有4*4 = 16个，如下图：

b\a 1   2   3   4

2   3   4   5   6

3   4   5   6   7

4   5   6   7   8

5   6   7   8   9

依次排序为：3,4,4,5, 5,5,6,6， 6,6，7,7, 7,8,8,9 （共16个）

解题思路：

最小堆方法实现的代码，其思路如下：

首先把a0+b0的结果放入堆中，此时堆中只有一个元素，自然满足最小堆条件，然后开始出堆的操作，从堆里面取出根节点（也就是最小的值），例如是a[i]+b[j]，则需要像最小堆中压入a[i+1]b[j] 和 a[i]+b[j+1]，当然，要保证下标不越界，如果下标越界了则忽略，另外要保证已经压入过堆中的组合（即使已经从堆中被取出了的）不再被压入堆中。不段进行出堆、入堆的操作，重复K次，就得到了K个最小的组合值。

堆的最大深度为logK,所以时间复杂度为K*logK数量级。

实现代码如下：

#include <iostream>
#include <vector>
#include <algorithm>
#include <map>
using namespace std;

struct Node{
	int posA;
	int posB;
	int sum;
	Node(){}
	Node(int a, int b, int s) :posA(a), posB(b), sum(s){}
};

class Solution{
public:
	static bool cmp(const Node& node1, const Node& node2)
	{
		return node1.sum > node2.sum;
	}
	vector<int> topKOfTwoSortedArray(const vector<int>& A, const vector<int>& B, int k)
	{
		vector<int> ans;
		if ((A.size() == 0 && B.size() == 0) || k<=0) return ans;
		int nA = A.size();
		int nB = B.size();
		if (k > nA*nB) k = nA*nB;
		map<pair<int,int>, bool> stats;
		vector<Node> heapAux;
		heapAux.push_back(Node(0,0,A[0]+B[0]));
		stats[make_pair(0,0)] = true;
		make_heap(heapAux.begin(), heapAux.end(), cmp);
		while (k>0){
			pop_heap(heapAux.begin(), heapAux.end(), cmp);
			Node tmp = heapAux.back();
			heapAux.pop_back();
			ans.push_back(tmp.sum);
			--k;
			int idxA = tmp.posA;
			int idxB = tmp.posB;
			if (idxA + 1 < nA && !stats[make_pair(idxA + 1, idxB)]){
				stats[make_pair(idxA + 1, idxB)] = true;
				heapAux.push_back(Node(idxA + 1, idxB, A[idxA + 1] + B[idxB]));
				push_heap(heapAux.begin(), heapAux.end(), cmp);				
			}
			if (idxB + 1 < nB && !stats[make_pair(idxA, idxB + 1)]){
				stats[make_pair(idxA, idxB + 1)] = true;
				heapAux.push_back(Node(idxA, idxB + 1, A[idxA] + B[idxB + 1]));
				push_heap(heapAux.begin(), heapAux.end(), cmp);
			}
		}
		return ans;
	}
};

int main()
{
	vector<int> A = {1,2,3,4};
	vector<int> B = {2,3,4,5};
	int k = 12;
	Solution sln;
	vector<int> ans = sln.topKOfTwoSortedArray(A, B, k);
	for (int i = 0; i < k; ++i){
		cout << ans[i] << " ";
	}
	cout << endl;
	return 0;
}

点赞