【LeetCode-373】找出第 k 小的距离对

问题

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

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

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

示例

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

解答

class Solution {
public:
    typedef tuple<int, int, int> TI;
    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        vector<vector<int>> res;
        priority_queue<TI, vector<TI>, greater<TI>> pq;
        auto push = [&](int i, int j) {
            if (i < nums1.size() && j < nums2.size())
                pq.emplace(nums1[i] + nums2[j], i, j);
        };
        push(0, 0);
        while (!pq.empty() && res.size() < k) {
            auto [_, i, j] = pq.top(); pq.pop();
            res.push_back({nums1[i], nums2[j]});
            push(i + 1, j);
            if (i == 0) push(i, j + 1);
        }
        return res;
    }
};

重点思路

需要注意的一个条件是，本题的两个数组已经排好序了。对于这类Top K问题，我们维护一个优先队列，重点在于需要放哪些数对入队。由上图可知，对于这类排好序的数组，我们只需要添加当前位置的右方数字，对于第一列的数字，我们还需要添加其下方位置的数字。