LeetCode-M373-堆-查找和最小的K对数字

题目

给定两个以升序排列的整形数组 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]

思路

遇到求第k大,第k小,前k大,前k小,就要想到优先对列,即最大堆或最小堆。求前k小,就可以通过维护一个k容量的最大堆,遍历完后堆中就是当前数据流中最小的k个元素;反之,求前k小,就维护一个k容量的最小堆,遍历完后堆中就是当前数据量中最大的k个元素。而对于第k大(第k小),也可以通过维护一个k容量的最小堆最大堆),遍历完后,堆顶即为第k大(小)元素。

思路1:这一题求前k小,即可以用最大堆来完成。但对于此题该解法存在一个效率问题,我们相当于暴力的将两个数组组合的所有情况压入堆中,会造成时间复杂度大大提高。
思路2:重新审查题目,可以看到题目给出了一个有效的前提,即每个数组递增排列。所以针对此情况,我们换个角度使用优先队列。
如果我们当前弹出的值是nums1[i],nums2[j], 那么我们下一个要添加的值应该是 num1[i],num2[j+1]。

如nums1 = {1,7,11,16} nums2 = {2,9,10,15},遍历nums1,分别与nums2的每一位元素组合,可以表示为:
(1,2) -> (1,9) -> (1,10) -> (1,15)
(7,2) -> (7,9) -> (7,10) -> (7,15)
(11,2) -> (11,9) -> (11,10) -> (11,15)
(16,2) -> (16,9) -> (16,10) -> (16,15)
针对上面这种表示方法,可以首先把一个列入堆。弹出一个后,下一个加入堆的是当前链表的下一个数字。

算法如下

  • 维护一个最小堆;
  • 先将第一列压入堆;
  • 取出堆顶第一个元素nums1[i],nums2[j],然后将nums1[i],nums2[j+1]压入堆;
  • 循环第二步,直到取出k个值或者当前堆为空。

解答

暴力解法

class Solution {
public:
    typedef vector vec;
    struct cmp{  //自定义最大堆比较方式
        bool operator()(const vec &left, const vec &right) const {
            return (left[0]+left[1]) < (right[0]+right[1]);
        }
    };
public:
    vector kSmallestPairs(vec& nums1, vec& nums2, int k) {
        if(nums1.empty() || nums2.empty()) return {};
        //存在k大于可组合的最大数量
        k = k < nums1.size()*nums2.size() ? k : nums1.size()*nums2.size();
        vec temp(2,0);
        vector ans(k,temp);
        priority_queue, cmp> que;
        for(auto i : nums1) {  //暴力将所有情况压入堆
            for(auto j : nums2){
                temp[0] = i;
                temp[1] = j;
                que.push(temp);   
                if(que.size() > k) que.pop();
            }
        }
        while(k > 0){  //倒序输出
            ans[k-1] = que.top();
            que.pop();
            k--;
        }
        return ans;
    }
};

索引法

class Solution {
public:
    typedef vector vec;
    struct data {  
        int num1;  //nums1的索引
        int num2;  //nums2的索引
        int sum;   //保存两个数组中元素的和
        data(int n1, int n2, int s): num1(n1), num2(n2), sum(s) {}
    };
    struct cmp { //自定义最小堆比较方式
        bool operator()(const data& left, const data& right) const { 
            return left.sum > right.sum;
        }
    };
public:
    vector kSmallestPairs(vec& nums1, vec& nums2, int k) {
        int len1 = nums1.size();
        int len2 = nums2.size();
        if(len1 < 1 || len2 < 1) return {};

        vector ans;
        priority_queue, cmp> que; //自定义最小堆
        for(int i=0; i 0) {
            data top = que.top();  //保存堆顶
            //首先将nums1与nums2中第一个元素输出
            ans.push_back({nums1[top.num1], nums2[top.num2]});
            que.pop();
            //若当前堆顶中第二个数组的索引未到最后,
            //就将后面一个元素重新压入堆中
            if(top.num2 < len2-1) 
                que.push(data(top.num1, top.num2+1, 
                              nums1[top.num1] + nums2[top.num2+1]));
            k--;
        }
        
        return ans;
    }
};

【关注公众号DoCode,每日一道LeetCode,将零碎时间利用起来】

你可能感兴趣的:(LeetCode-M373-堆-查找和最小的K对数字)