373. 查找和最小的K对数字

题目：

链接：https://leetcode-cn.com/problems/find-k-pairs-with-smallest-sums/

给定两个以升序排列的整形数组 nums1 和 nums2, 以及一个整数 k。

定义一对值 (u,v)，其中第一个元素来自 nums1，第二个元素来自 nums2。

找到和最小的 k 对数字 (u1,v1), (u2,v2) ... (uk,vk)。

示例 1:

输入: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
输出: [1,2],[1,4],[1,6]
解释: 返回序列中的前 3 对数：
[1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]
示例 2:

输入: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
输出: [1,1],[1,1]
解释: 返回序列中的前 2 对数：
  [1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]
示例 3:

输入: nums1 = [1,2], nums2 = [3], k = 3
输出: [1,3],[2,3]
解释: 也可能序列中所有的数对都被返回:[1,3],[2,3]

 

解答：

方法1：

直接套用TOP-K问题的解法（https://www.cnblogs.com/FdWzy/p/12321325.html）

由于求最小的k个元素对，那么首先定义小于重载关系。

之后利用最大堆，堆大小不超过k直接push，超过k要看是否比堆顶小，是则pop并push，否则跳过。

效率O(m*n*log k)

 

class Solution {
public:
    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        vector<vector<int>> res;
        int len1=nums1.size(),len2=nums2.size();
        auto cmp=[](vector<int>& a,vector<int>& b){return a[0]+a[1]<b[0]+b[1];};
        priority_queue<vector<int>,vector<vector<int>>,decltype(cmp)> max_heap(cmp);
        vector<int> tmp1,tmp2;
        for(int i=0;i<len1;++i){
            for(int j=0;j<len2;++j){
                if(max_heap.size()<k){
                    max_heap.push({nums1[i],nums2[j]});
                }
                else{
                    tmp1=max_heap.top();
                    tmp2={nums1[i],nums2[j]};
                    if(not cmp(tmp1,tmp2)){
                        max_heap.pop();
                        max_heap.push(move(tmp2));
                    }
                }
            }
        }
        while(not max_heap.empty()){
            res.emplace_back(max_heap.top());
            max_heap.pop();
        }
        reverse(res.begin(),res.end());
        return res;
    }
};

 

 

 

 

方法2：

前一个方法比较暴力，直接枚举所有可能的元素对全部加入最大堆进行筛选。这样呢，没有利用两个数组有序的性质，两个无序数组也可以这样做。

所以我们考虑记录对于每一个nums1中的索引i，如果已经把[i,0],[i,1],[i,2]...[i,j]全部的元素对加入了结果，那么下次再查找[i,xx]元素对的时候，就可以从j+1开始找，因为数组是有序的。

 

用一个memo数组记录。

做法，我们这回用最小堆，堆顶直接指示为最小的元素对。每次取堆顶，更改相应的memo数组值。如果memo[i]=len(nums2)，那说明[i,xx]元素对已经无法再取了，就直接从堆中删掉。

如果memo[i]<len(nums2)，那么再push回堆里。

效率：O(m*log m+k*log m)

class Solution {
public:
    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        vector<vector<int>> res;
        int len1=nums1.size(),len2=nums2.size();
        if(len1==0 or len2==0){return res;}
        vector<int> memo(len1,-1);
        //memo[i]指示：{nums1[i],nums2[j]}元素对中上次已经放入结果数组那个的j的值，i对应第一个未使用的j为memo[i]+1
        auto cmp=[&](int x,int y){return nums1[x]+nums2[memo[x]+1]>nums1[y]+nums2[memo[y]+1];};
        priority_queue<int,vector<int>,decltype(cmp)> min_heap(cmp);//小顶堆
        for(int i=0;i<len1;++i){
            min_heap.push(i);
        }
        while(res.size()<k and not min_heap.empty()){//每次循环把小顶堆的堆顶（也就是目前的最小对）push进结果
            int tp=min_heap.top();
            min_heap.pop();
            res.push_back({nums1[tp],nums2[memo[tp]+1]});
            memo[tp]+=1;                            //memo[tp]+1已经遍历，所以memo[tp]+=1；
            if(memo[tp]<len2-1){                    //如果tp对应nums2中的索引到达nums2尾部，那就不push了
                min_heap.push(tp);
            }
        }
        return res;
    }
};