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

优先队列

• 思路：
  • 使用下标 (x, y) 标识数值对，x 为第一个数组的下标，y 为第二个数组的下标；
  • 所以 k 个数值对 x 的范围属于 [0, min(k, m)]，m 为第一个数组的 size；
  • 数值对 (x, y) ，那么下一个比其大的数组对是 min{(x, y + 1), (x + 1, y)}；
  • 可以先固定 x ，即将 x 可能的值全选，来动态变更 y；
  • 构建一个优先队列，存放的是 (x, y)，小顶堆，即其对应的数值对的和最小的总是在堆顶：nums1[a.first] + nums2[a.second] > nums1[b.first] + nums2[b.second];
  • 将小顶堆取 k 次堆顶即可，每次之后将 (x, y + 1) 入堆即可；

class Solution {
public:
    vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
        auto cmp = [&nums1, &nums2](const std::pair<int, int>& a, const std::pair<int, int>& b) {
            return nums1[a.first] + nums2[a.second] > nums1[b.first] + nums2[b.second];
        };

        int m = nums1.size();
        int n = nums2.size();

        std::vector<std::vector<int>> result;
        std::priority_queue<std::pair<int, int>, std::vector<std::pair<int, int>>, decltype(cmp)> pq(cmp);

        for (int i = 0; i < std::min(k, m); i++) {
            pq.emplace(i, 0);
        }

        while (k-- > 0 && !pq.empty()) {
            auto [x, y] = pq.top();
            pq.pop();
            result.push_back(std::initializer_list<int>{nums1[x], nums2[y]});

            if (y + 1 < n) {
                pq.emplace(x, y + 1);
            }
        }
        
        return result;
    }
};