[leetcode] 373. Find K Pairs with Smallest Sums

You are given two integer arrays nums1 and nums2 sorted in ascending order and an integer k.

Define a pair (u,v) which consists of one element from the first array and one element from the second array.

Find the k pairs (u1,v1),(u2,v2) ...(uk,vk) with the smallest sums.

Example 1:

Given nums1 = [1,7,11], nums2 = [2,4,6],  k = 3

Return: [1,2],[1,4],[1,6]

The first 3 pairs are returned from the sequence:
[1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]

Example 2:

Given nums1 = [1,1,2], nums2 = [1,2,3],  k = 2

Return: [1,1],[1,1]

The first 2 pairs are returned from the sequence:
[1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]

Example 3:

Given nums1 = [1,2], nums2 = [3],  k = 3 

Return: [1,3],[2,3]

All possible pairs are returned from the sequence:
[1,3],[2,3]

这道题是找两个数组中相加和最小的K对数字，题目难度为Medium。

比较直观的想法是用堆来解决问题，如果用小根堆需要存入堆的数据较多，复杂度相应也较高，这里用大根堆来存储数字对，相加和小于堆顶数字对的将其存入堆中，如果堆中元素个数大于K，删除堆顶元素，最后将整个堆中元素存入结果。具体代码：

class Solution {
    typedef pair<int, int> data;
    
    struct cmp {
        bool operator()(data& l, data& r) {
            return l.first + l.second < r.first + r.second;
        }
    };
public:
    vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        priority_queue<data, vector<data>, cmp> heap;
        vector<data> ret;
        int m = min((int)nums1.size(), k);
    ​    ​int n = min((int)nums2.size(), k);
        
        for(int i=0; i<m; ++i) {
            for(int j=0; j<n; ++j) {
                if(heap.size() < k) {
                    heap.push(make_pair(nums1[i], nums2[j]));
                }
                else if(nums1[i]+nums2[j] < heap.top().first + heap.top().second) {
                    heap.push(make_pair(nums1[i], nums2[j]));
                    heap.pop();
                }
            }
        }
        
        while(!heap.empty()) {
            ret.push_back(heap.top());
            heap.pop();
        }
        
        return ret;
    }
};

两个数组都已排序，我们可以依次选出当前最小的数字对，记录下nums1中每个数字关联的数字在nums2中取到了哪里，这样遍历nums1即可找出当前相加和最小的数字对，将其存入结果并将nums1中选定数字关联的nums2中位置向前进1，经过K此选择即可得到最终结果。具体代码：

class Solution {
public:
    vector<pair<int, int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        vector<pair<int, int>> ret;
        int m = nums1.size(), n = nums2.size();
        int cnt = min(k, m*n);
        vector<int> idx(m, 0);
        
        for(int i=0; i<cnt; ++i) {
            int curIdx = 0, curMin = INT_MAX;
            for(int j=0; j<m; ++j) {
                if(idx[j]<n && nums1[j]+nums2[idx[j]]<curMin) {
                    curIdx = j;
                    curMin = nums1[j] + nums2[idx[j]];
                }
            }
            ret.push_back(make_pair(nums1[curIdx], nums2[idx[curIdx]]));
            ++idx[curIdx];
        }
        
        return ret;
    }
};