查找和最小的k对数字

10 Sep 2024 in Algorithm Pageviews

leetcode!

给定两个以 非递减顺序排列 的整数数组 nums₁ 和 nums₂ , 以及一个整数 k 。定义一对值 (u,v)，其中第一个元素来自 nums₁，第二个元素来自 nums₂ 。请找到和最小的 k 个数对 (u₁, v₁), (u₂, v₂) … (u_k, v_k) 。

示例 1:

输入: nums₁ = [1,7,11], nums₂ = [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:

输入: nums₁ = [1, 1, 2], nums₂ = [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]

代码实现如下：

#include <bits/stdc++.h>

using namespace std;

class Solution {
public:
    vector<vector<int>> KSmallestPairs(vector<int> &nums1, vector<int> &nums2, int k)
    {
        auto Com = [&nums1, &nums2](const pair<int, int> &a, const pair<int, int> &b) {
            return make_tuple(-nums1[a.first] - nums2[a.second]) <
                   make_tuple(-nums1[b.first] - nums2[b.second]);
        };
        priority_queue<pair<int, int>, vector<pair<int, int>>, decltype(Com)> pq(Com);
        for (size_t i = 0; i < nums1.size(); i++) {
           pq.push({i, 0});
        }
        vector<vector<int>> res;
        while ((k--) && !pq.empty()) {
            auto [i, j] = pq.top();
            pq.pop();
            res.emplace_back(vector<int>({nums1[i], nums2[j]}));
            if (j + 1 < nums2.size()) {
                pq.push({i, j + 1});
            }
        }
        return res;
    }
};