力扣 373. 查找和最小的 K 对数字 [堆]

373. 查找和最小的 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]

提示:

• 1 <= nums1.length, nums2.length <= 105
• -109 <= nums1[i], nums2[i] <= 109
• nums1 和 nums2 均为升序排列
• 1 <= k <= 1000

题解

思路来自官方。

首先需要了解一下堆，可以查看堆介绍与实现。

下标选择

要求得到一对值u,v，分别来自题目中的两个升序排列数组，为了方便访问每种情况，用i，j分别表示nums1和nums2元素的下标，在存入结果ans时，通过访问数组来获取下标对应的值u,v.

利用升序数组的性质，第一小的和值对为nums1[0]+nums2[0]，而最大的和值对为nums1[len-1]+nums2[len-2]。

因为当前下标对为i,j时，那么比当前和大且最接近当前和的新下标对为i+1,j或者i,j+1。

假设当前已选的前n个对的下标为(a1,b2),(a2,b2)...(an,bn)，根据升序的性质，那么第n+1个对可以从前n个对中寻找，即第n+1个对的下标范围是：

(a1+1,b1),(a1,b1+1)...(an+1,bn),(an,bn+1)，将第n+1个下标对的范围全部入堆（优先队列实现），由最小堆动态排序。

假设我们利用堆的特性可以求出待选范围中最小数对的索引为 i,j，同时将新的待选的下标对 i+1,j和i,j+1 加入到堆中，直到我们选出 k 个数对即可。

下标去重

直接将待选下标i+1,j和i,j+1 加入到堆中则可能出现重复的问题，一般需要设置标记位解决去重的问题。我们可以将nums1 的前 k 个下标对 (0,0),(1,0),(k-1,0)(0,0),(1,0),…,(k−1,0) 加入到队列中，每次从队列中取出元素 (x,y)时，我们只需要将nums2 中的下标增加即可，这样避免了重复加入元素的问题。

也就是先将和第一小的下标对的待选范围放在 nums1[i]，nums2[0]到 nums1[len1]，nums2[0]中，每次只右移num2中的下标j，来获取待选的下标对，而nums1中的下标没有变化过。

查看代码

class Solution {
public:
    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        //排序方法：若前一个pair的两个下标得到的元素之和>第二个pair得到的元素之和，顺序不变，否则交换,即按和的从大到小排列，得到小顶堆，top为最小值
        auto cmp = [&nums1, &nums2](const pair<int, int> & a, const pair<int, int> & b) {
            return nums1[a.first] + nums2[a.second] > nums1[b.first] + nums2[b.second];
        };

        int len1 = nums1.size();
        int len2 = nums2.size();
        vector<vector<int>> ans;   
        priority_queue<pair<int, int>, vector<pair<int, int>>, decltype(cmp)> pq(cmp);
        for (int i = 0; i < min(k, len1); i++) {//num1中每个元素分别与num2中第一个元素组成下标对
            pq.emplace(i, 0);
        }
        while (k-- > 0 && !pq.empty()) {
            auto [x, y] = pq.top(); //按和从小到大顺序获取下标对，假如当前为第m小的
            pq.pop();
            ans.emplace_back(initializer_list<int>{nums1[x], nums2[y]});//initializer_listinitializer_list是C++11中新增的一个模板类，在初始化时调用，在使用STL的时候，通过{}对容器进行初始化就会调用
            //接下来num2的下标右移，入堆，堆会自动排序，这样下一轮循环的时候，top会得到m+1小的
            if (y + 1 < len2) {
                pq.emplace(x, y + 1);
            }
        }
        return ans;
    }
};