面试经典150题 P373 查找和最小的K对数字

---

时间轴

时间轴

2025-12-08

init

堆

---

题目：

P373 查找和最小的K对数字

https://leetcode.cn/problems/find-k-pairs-with-smallest-sums/description/?envType=study-plan-v2&envId=top-interview-150

这题不能用双指针，因为不是线性的，用堆解决的核心是先把 {nums1[0],nums2[j]} 都放入堆中(0 <= j <= nums2.size()),然后从堆中去除最小和的那一对{i,j}, j 不变的情况下比它大的值中{i+1, j}最小。

那为什么可以让 j 不变呢？因为 j 的推进（即 j+1）已经在初始化阶段全部加入堆({num1[0], nums2[j]})

示例 1 的nums1 = [1, 7, 11]，nums2 = [2, 4, 6]。我们把每个数对的和算出来，可以得到一个矩阵 $M$，其中 $M_{i,j} = nums1[i] + nums2[j]$。

$$M = \begin{bmatrix} 3 & 5 & 7 \\ 9 & 11 & 13 \\ 13 & 15 & 17 \end{bmatrix}$$

由于nums2是递增的，所以矩阵每一行都是递增的。问题相当于： 合并 $n$ 个升序列表，找前 $k$ 小元素。（其中 $n$ 是nums1的长度） 根据 P23 合并 K 个升序链表 的堆的做法：

1. 把矩阵每一行的第一个数 $M_{i,0}$ 及其位置 $(i, 0)$ 加到最小堆中。
2. 循环 $k$ 次。
3. 每次循环，弹出堆顶，把堆顶 $M*{i,j}$ 的对应数对加入答案，把堆顶右边元素 $M*{i,j+1}$ 及其位置 $(i, j+1)$ 入堆。

#include <vector>
#include <queue>
#include <utility>
#include <array>
using std::vector;
using std::pair;
using std::priority_queue;
using std::array;

class Solution {
    public:
        vector<vector<int> > kSmallestPairs(vector<int> &nums1,
                                            vector<int> &nums2, int k)
        {
                vector<vector<int> > res;
                if (nums1.empty() || nums2.empty() || k == 0)
                        return res;

                auto cmp = [&](const array<int, 2> &a, const array<int, 2> &b) {
                        return nums1[a[0]] + nums2[a[1]] >
                               nums1[b[0]] + nums2[b[1]];
                };
                priority_queue<array<int, 2>, vector<array<int, 2> >,
                               decltype(cmp)>
                        pq(cmp);


                for (int i = 0; i < nums1.size() && i < k; i++) {
                        pq.push({ i, 0 });
                }

                while (k-- && !pq.empty()) {
                        auto [i, j] = pq.top();
                        pq.pop();
                        res.push_back({ nums1[i], nums2[j] });

                        if (j + 1 < nums2.size()) {
                                pq.push({ i, j+1 });
                        }
                }

                return res;
        }
};