373. 查找和最小的 K 对数字
题目描述
输入两个按非递减顺序排好的整数数组 nums1、nums2，以及一个整数 k。
从 nums1 取一个数 u、从 nums2 取一个数 v，组成数对 [u, v]。不同位置上的元素视为不同的选择，所以值相同的数对可能出现多次。
在所有可能的数对中，找出 u + v 最小的 k 个，按和从小到大的顺序返回。和相同的数对之间先后不限。
示例 1
输入：nums1 = [1,7,11], nums2 = [2,4,6], k = 3 输出：[[1,2],[1,4],[1,6]]
解释：和最小的三对是 [1,2]、[1,4]、[1,6]，和分别为 3、5、7。
示例 2
输入：nums1 = [1,1,2], nums2 = [1,2,3], k = 2 输出：[[1,1],[1,1]]
解释：nums1 的两个 1 分别和 nums2 的 1 配对，得到两个和为 2 的数对。
数据范围
1 <= nums1.length, nums2.length <= 10^5-10^9 <= nums1[i], nums2[i] <= 10^9nums1和nums2都已按升序排列1 <= k <= 10^4k <= nums1.length * nums2.length
解题思路：最小堆多路归并
把所有数对看成一张表：第 i 行是 nums1[i] 搭配 nums2 的每个数，每一行从左到右和递增。这就像合并 n 条有序链表：先把每行的第一个数对 (i,0) 放进最小堆，每次弹出和最小的 (i,j) 放入答案，再把同一行的下一个 (i,j+1) 推入堆。行数只需取前 min(k, len(nums1)) 行，因为更靠下的行开头更大，前 k 名不可能全落在后面。
为什么正确
堆里每一行最多只有一个候选，而且是该行还没输出的最小数对；所以堆顶就是全部未输出数对中和最小的那个。
复杂度分析
时间 O(k log min(k, m))，空间 O(min(k, m))，m 为 nums1 的长度。
易错点
不要把 (i+1,j) 和 (i,j+1) 都推进堆，那样同一个数对会被重复推入，需要额外去重。只沿行向右推进，初始时放好每行开头即可。另外两个数之和可能超过 32 位整数范围。
参考代码
Python、Java、C++、Go、JavaScript、TypeScript 六种语言，逻辑一致。
class Solution:
def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
heap = [(nums1[i] + nums2[0], i, 0) for i in range(min(k, len(nums1)))]
heapq.heapify(heap)
ans = []
while heap and len(ans) < k:
s, i, j = heapq.heappop(heap)
ans.append([nums1[i], nums2[j]])
if j + 1 < len(nums2):
heapq.heappush(heap, (nums1[i] + nums2[j + 1], i, j + 1))
return ansclass Solution {
public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
PriorityQueue<long[]> heap = new PriorityQueue<>((a, b) -> a[0] != b[0] ? Long.compare(a[0], b[0]) : a[1] != b[1] ? Long.compare(a[1], b[1]) : Long.compare(a[2], b[2]));
for (int i = 0; i < Math.min(k, nums1.length); i++) heap.offer(new long[]{(long) nums1[i] + nums2[0], i, 0});
List<List<Integer>> ans = new ArrayList<>();
while (!heap.isEmpty() && ans.size() < k) {
long[] t = heap.poll();
int i = (int) t[1], j = (int) t[2];
ans.add(Arrays.asList(nums1[i], nums2[j]));
if (j + 1 < nums2.length) {
heap.offer(new long[]{(long) nums1[i] + nums2[j + 1], i, j + 1});
}
}
return ans;
}
}class Solution {
public:
vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
priority_queue<tuple<long long, int, int>, vector<tuple<long long, int, int>>, greater<>> heap;
for (int i = 0; i < min(k, (int)nums1.size()); i++) heap.emplace((long long)nums1[i] + nums2[0], i, 0);
vector<vector<int>> ans;
while (!heap.empty() && (int)ans.size() < k) {
auto [s, i, j] = heap.top();
heap.pop();
ans.push_back({nums1[i], nums2[j]});
if (j + 1 < (int)nums2.size()) {
heap.emplace((long long)nums1[i] + nums2[j + 1], i, j + 1);
}
}
return ans;
}
};type pairHeap [][3]int
func (h pairHeap) Len() int { return len(h) }
func (h pairHeap) Less(a, b int) bool {
for t := 0; t < 3; t++ {
if h[a][t] != h[b][t] {
return h[a][t] < h[b][t]
}
}
return false
}
func (h pairHeap) Swap(a, b int) { h[a], h[b] = h[b], h[a] }
func (h *pairHeap) Push(x any) { *h = append(*h, x.([3]int)) }
func (h *pairHeap) Pop() any {
old := *h
x := old[len(old)-1]
*h = old[:len(old)-1]
return x
}
func kSmallestPairs(nums1 []int, nums2 []int, k int) [][]int {
h := &pairHeap{}
for i := 0; i < min(k, len(nums1)); i++ {
*h = append(*h, [3]int{nums1[i] + nums2[0], i, 0})
}
heap.Init(h)
ans := [][]int{}
for h.Len() > 0 && len(ans) < k {
t := heap.Pop(h).([3]int)
i, j := t[1], t[2]
ans = append(ans, []int{nums1[i], nums2[j]})
if j+1 < len(nums2) {
heap.Push(h, [3]int{nums1[i] + nums2[j+1], i, j + 1})
}
}
return ans
}class MinHeap {
constructor(less) { this.a = []; this.less = less; }
get size() { return this.a.length; }
push(x) {
const a = this.a; a.push(x);
for (let i = a.length - 1; i > 0;) {
const p = (i - 1) >> 1;
if (!this.less(a[i], a[p])) break;
[a[i], a[p]] = [a[p], a[i]]; i = p;
}
}
pop() {
const a = this.a, top = a[0], last = a.pop();
if (a.length) {
a[0] = last;
for (let i = 0; ;) {
let m = i; const l = 2 * i + 1, r = l + 1;
if (l < a.length && this.less(a[l], a[m])) m = l;
if (r < a.length && this.less(a[r], a[m])) m = r;
if (m === i) break;
[a[i], a[m]] = [a[m], a[i]]; i = m;
}
}
return top;
}
}
var kSmallestPairs = function(nums1, nums2, k) {
const less = (a, b) => a[0] !== b[0] ? a[0] < b[0] : a[1] !== b[1] ? a[1] < b[1] : a[2] < b[2];
const heap = new MinHeap(less);
for (let i = 0; i < Math.min(k, nums1.length); i++) heap.push([nums1[i] + nums2[0], i, 0]);
const ans = [];
while (heap.size && ans.length < k) {
const [s, i, j] = heap.pop();
ans.push([nums1[i], nums2[j]]);
if (j + 1 < nums2.length) {
heap.push([nums1[i] + nums2[j + 1], i, j + 1]);
}
}
return ans;
};class MinHeap<T> {
private a: T[] = [];
constructor(private less: (x: T, y: T) => boolean) {}
get size(): number { return this.a.length; }
push(x: T): void {
const a = this.a; a.push(x);
for (let i = a.length - 1; i > 0;) {
const p = (i - 1) >> 1;
if (!this.less(a[i], a[p])) break;
[a[i], a[p]] = [a[p], a[i]]; i = p;
}
}
pop(): T {
const a = this.a, top = a[0], last = a.pop()!;
if (a.length) {
a[0] = last;
for (let i = 0; ;) {
let m = i; const l = 2 * i + 1, r = l + 1;
if (l < a.length && this.less(a[l], a[m])) m = l;
if (r < a.length && this.less(a[r], a[m])) m = r;
if (m === i) break;
[a[i], a[m]] = [a[m], a[i]]; i = m;
}
}
return top;
}
}
function kSmallestPairs(nums1: number[], nums2: number[], k: number): number[][] {
const less = (a: number[], b: number[]) => a[0] !== b[0] ? a[0] < b[0] : a[1] !== b[1] ? a[1] < b[1] : a[2] < b[2];
const heap = new MinHeap<number[]>(less);
for (let i = 0; i < Math.min(k, nums1.length); i++) heap.push([nums1[i] + nums2[0], i, 0]);
const ans: number[][] = [];
while (heap.size && ans.length < k) {
const [s, i, j] = heap.pop();
ans.push([nums1[i], nums2[j]]);
if (j + 1 < nums2.length) {
heap.push([nums1[i] + nums2[j + 1], i, j + 1]);
}
}
return ans;
}其他解法
最小堆按和从小到大逐个产出数对，适合 k 较小；二分查找先确定第 k 小的和再一次性收集，空间更省但结果无序。
解法 2：二分查找第 k 小的和 时间 O((m + n) log D + k)，空间 O(1)
在数对和的取值范围上二分，找最小的 s，使和不超过 s 的数对至少有 k 个；计数时 i 递增、j 从 nums2 末尾递减，一次扫描完成。然后先收集所有和小于 s 的数对，再补和等于 s 的数对，直到凑满 k 个。D 为最大和与最小和之差。
不需要堆，额外空间 O(1)，k 很大时更稳定。输出不按和排序，只保证是和最小的 k 个；代码比堆更长，边界更容易写错。
class Solution:
def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
m, n = len(nums1), len(nums2)
def count(x):
cnt, j = 0, n - 1
for i in range(m):
while j >= 0 and nums1[i] + nums2[j] > x:
j -= 1
cnt += j + 1
if cnt >= k:
break
return cnt
lo, hi = nums1[0] + nums2[0], nums1[-1] + nums2[-1]
while lo < hi:
mid = (lo + hi) // 2
if count(mid) >= k:
hi = mid
else:
lo = mid + 1
s = lo
ans = []
j = n - 1
for i in range(m):
while j >= 0 and nums1[i] + nums2[j] >= s:
j -= 1
for t in range(j + 1):
ans.append([nums1[i], nums2[t]])
j = n - 1
for i in range(m):
if len(ans) >= k:
break
while j >= 0 and nums1[i] + nums2[j] > s:
j -= 1
t = j
while t >= 0 and nums1[i] + nums2[t] == s and len(ans) < k:
ans.append([nums1[i], nums2[t]])
t -= 1
return ansclass Solution {
public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
int m = nums1.length, n = nums2.length;
long lo = (long) nums1[0] + nums2[0], hi = (long) nums1[m - 1] + nums2[n - 1];
while (lo < hi) {
long mid = (lo + hi) >> 1;
if (count(nums1, nums2, mid, k) >= k) hi = mid;
else lo = mid + 1;
}
long s = lo;
List<List<Integer>> ans = new ArrayList<>();
int j = n - 1;
for (int i = 0; i < m; i++) {
while (j >= 0 && (long) nums1[i] + nums2[j] >= s) j--;
for (int t = 0; t <= j; t++) ans.add(Arrays.asList(nums1[i], nums2[t]));
}
j = n - 1;
for (int i = 0; i < m && ans.size() < k; i++) {
while (j >= 0 && (long) nums1[i] + nums2[j] > s) j--;
for (int t = j; t >= 0 && (long) nums1[i] + nums2[t] == s && ans.size() < k; t--) {
ans.add(Arrays.asList(nums1[i], nums2[t]));
}
}
return ans;
}
private long count(int[] nums1, int[] nums2, long x, int k) {
long cnt = 0;
int j = nums2.length - 1;
for (int i = 0; i < nums1.length; i++) {
while (j >= 0 && (long) nums1[i] + nums2[j] > x) j--;
cnt += j + 1;
if (cnt >= k) break;
}
return cnt;
}
}class Solution {
public:
vector<vector<int>> kSmallestPairs(vector<int>& nums1, vector<int>& nums2, int k) {
int m = nums1.size(), n = nums2.size();
auto count = [&](long long x) {
long long cnt = 0;
int j = n - 1;
for (int i = 0; i < m; i++) {
while (j >= 0 && (long long)nums1[i] + nums2[j] > x) j--;
cnt += j + 1;
if (cnt >= k) break;
}
return cnt;
};
long long lo = (long long)nums1[0] + nums2[0], hi = (long long)nums1[m - 1] + nums2[n - 1];
while (lo < hi) {
long long mid = (lo + hi) >> 1;
if (count(mid) >= k) hi = mid;
else lo = mid + 1;
}
long long s = lo;
vector<vector<int>> ans;
int j = n - 1;
for (int i = 0; i < m; i++) {
while (j >= 0 && (long long)nums1[i] + nums2[j] >= s) j--;
for (int t = 0; t <= j; t++) ans.push_back({nums1[i], nums2[t]});
}
j = n - 1;
for (int i = 0; i < m && (int)ans.size() < k; i++) {
while (j >= 0 && (long long)nums1[i] + nums2[j] > s) j--;
for (int t = j; t >= 0 && (long long)nums1[i] + nums2[t] == s && (int)ans.size() < k; t--) {
ans.push_back({nums1[i], nums2[t]});
}
}
return ans;
}
};func kSmallestPairs(nums1 []int, nums2 []int, k int) [][]int {
m, n := len(nums1), len(nums2)
count := func(x int) int {
cnt, j := 0, n-1
for i := 0; i < m; i++ {
for j >= 0 && nums1[i]+nums2[j] > x {
j--
}
cnt += j + 1
if cnt >= k {
break
}
}
return cnt
}
lo, hi := nums1[0]+nums2[0], nums1[m-1]+nums2[n-1]
for lo < hi {
mid := (lo + hi) >> 1
if count(mid) >= k {
hi = mid
} else {
lo = mid + 1
}
}
s := lo
ans := [][]int{}
j := n - 1
for i := 0; i < m; i++ {
for j >= 0 && nums1[i]+nums2[j] >= s {
j--
}
for t := 0; t <= j; t++ {
ans = append(ans, []int{nums1[i], nums2[t]})
}
}
j = n - 1
for i := 0; i < m && len(ans) < k; i++ {
for j >= 0 && nums1[i]+nums2[j] > s {
j--
}
for t := j; t >= 0 && nums1[i]+nums2[t] == s && len(ans) < k; t-- {
ans = append(ans, []int{nums1[i], nums2[t]})
}
}
return ans
}var kSmallestPairs = function(nums1, nums2, k) {
const m = nums1.length, n = nums2.length;
const count = (x) => {
let cnt = 0, j = n - 1;
for (let i = 0; i < m; i++) {
while (j >= 0 && nums1[i] + nums2[j] > x) j--;
cnt += j + 1;
if (cnt >= k) break;
}
return cnt;
};
let lo = nums1[0] + nums2[0], hi = nums1[m - 1] + nums2[n - 1];
while (lo < hi) {
const mid = Math.floor((lo + hi) / 2);
if (count(mid) >= k) hi = mid;
else lo = mid + 1;
}
const s = lo;
const ans = [];
let j = n - 1;
for (let i = 0; i < m; i++) {
while (j >= 0 && nums1[i] + nums2[j] >= s) j--;
for (let t = 0; t <= j; t++) ans.push([nums1[i], nums2[t]]);
}
j = n - 1;
for (let i = 0; i < m && ans.length < k; i++) {
while (j >= 0 && nums1[i] + nums2[j] > s) j--;
for (let t = j; t >= 0 && nums1[i] + nums2[t] === s && ans.length < k; t--) {
ans.push([nums1[i], nums2[t]]);
}
}
return ans;
};function kSmallestPairs(nums1: number[], nums2: number[], k: number): number[][] {
const m = nums1.length, n = nums2.length;
const count = (x: number): number => {
let cnt = 0, j = n - 1;
for (let i = 0; i < m; i++) {
while (j >= 0 && nums1[i] + nums2[j] > x) j--;
cnt += j + 1;
if (cnt >= k) break;
}
return cnt;
};
let lo = nums1[0] + nums2[0], hi = nums1[m - 1] + nums2[n - 1];
while (lo < hi) {
const mid = Math.floor((lo + hi) / 2);
if (count(mid) >= k) hi = mid;
else lo = mid + 1;
}
const s = lo;
const ans: number[][] = [];
let j = n - 1;
for (let i = 0; i < m; i++) {
while (j >= 0 && nums1[i] + nums2[j] >= s) j--;
for (let t = 0; t <= j; t++) ans.push([nums1[i], nums2[t]]);
}
j = n - 1;
for (let i = 0; i < m && ans.length < k; i++) {
while (j >= 0 && nums1[i] + nums2[j] > s) j--;
for (let t = j; t >= 0 && nums1[i] + nums2[t] === s && ans.length < k; t--) {
ans.push([nums1[i], nums2[t]]);
}
}
return ans;
}