题目

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

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

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

来源：力扣（LeetCode）

解答

Tips：一般看到第K个、前K个等词语，就会想到用堆这个数据结构。

理解题意，就是nums1和nums2的所有组合，如果直接将所有组合加入到堆中进行排序，并取前K个组合，一定会超时。

题目给定了一个条件，就是nums1和nums2都是有序的数组，所以可以根据这一点简化算法复杂度。

可以将问题分解为一个n路归并的问题，拆解开来就是：将nums1的每个元素，都分配一个nums2数组（可见下图），然后就可以组成所有的组合排列；这里面最小的，一定是左侧的nums1的一个元素，与右侧nums2元素最小下标的值所组成的和。

所以：

1. 建立堆，堆的基本元素，包括一个和，还有nums1 & nums2的下标
2. 将【nums1的每个元素和nums2的第一个元素】组成的基本元素入堆
3. 循环出堆，出堆加入到返回值中，同时，如果出堆所对应的元素还有后续元素（即nums2的下标还没有到最后一个），就将后续元素入堆
4. 最终取完K组数组，或取完所有排列，就返回结果。

代码如下：

// 存放到堆中的基本元素单位
struct tripleNums {
    int sum; // 两数之和
    int firstIndex; // 存的是nums1的下标，而非数
    int secondIndex; // 同上，存的是nums2的下标

    bool operator<(const tripleNums &a) const {
        return sum > a.sum;
    }
};

class Solution {
public:
    vector<vector<int>> kSmallestPairs(vector<int> &nums1, vector<int> &nums2, int k) {
        vector<vector<int>> ret;

        int n1 = nums1.size();
        int n2 = nums2.size();

        priority_queue<tripleNums> pq;

        // 将 n1 个元素入堆
        for (int i = 0; i < n1; ++i) {
            pq.push(tripleNums{nums1[i] + nums2[0], i, 0});
        }

        // 每次从堆中取一个元素，一定是最小值，且取一个要再放一个（如果还有数的话）
        while (!pq.empty() && k--) {
            // 取小根堆的顶
            auto get = pq.top();
            pq.pop();
            // 将取出的元素加入返回值中
            ret.push_back(vector<int>{nums1[get.firstIndex], nums2[get.secondIndex]});

            // 如果该路元素还有后续元素，则入堆
            if (get.secondIndex < n2 - 1) {
                pq.push(tripleNums{nums1[get.firstIndex] + nums2[get.secondIndex + 1], get.firstIndex,
                                   get.secondIndex + 1});
            }
        }

        return ret;
    }
};