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

题目描述：

有两个大小分别是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;
}