373. Find K Pairs with Smallest Sums
373. Find K Pairs with Smallest Sums
题目
You are given two integer arrays nums1
and nums2
sorted in non- decreasing order and an integer k
.
Define a pair (u, v)
which consists of one element from the first array and one element from the second array.
Return the k
pairs (u1, v1), (u2, v2), ..., (uk, vk)
with the smallest sums.
Example 1:
Input: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
Output: [[1,2],[1,4],[1,6]]
Explanation: The first 3 pairs are returned from the sequence: [1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]
Example 2:
Input: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
Output: [[1,1],[1,1]]
Explanation: The first 2 pairs are returned from the sequence: [1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]
Example 3:
Input: nums1 = [1,2], nums2 = [3], k = 3
Output: [[1,3],[2,3]]
Explanation: All possible pairs are returned from the sequence: [1,3],[2,3]
Constraints:
1 <= nums1.length, nums2.length <= 10^5
-10^9 <= nums1[i], nums2[i] <= 10^9
nums1
andnums2
both are sorted in non-decreasing order.1 <= k <= 10^4
k <= nums1.length * nums2.length
题目大意
给定两个以 非递减顺序排列 的整数数组 nums1
和 nums2
, 以及一个整数 k
。
定义一对值 (u,v)
,其中第一个元素来自 nums1
,第二个元素来自 nums2
。
请找到和最小的 k
个数对 (u1,v1)
, (u2,v2)
... (uk,vk)
。
解题思路
可以通过优先队列(堆)来解决。我们可以维护一个小顶堆,堆中存储的是每一个可能的数对的和,同时记录这个数对在两个数组中的位置。每次弹出堆顶元素,并将下一个可能的数对入堆。最终返回的结果数组就是从小到大的前 k
个数对。
- 构建一个小顶堆,堆中的每个元素是一个三元组
(sum, i, j)
,其中sum
表示nums1[i] + nums2[j]
的和,i
和j
分别表示这个和在两个数组中的位置。 - 初始化堆中元素,将
(nums1[0] + nums2[0], 0, 0)
加入堆中。 - 重复以下步骤
k
次:- 弹出堆顶元素
(sum, i, j)
,将(nums1[i+1] + nums2[j], i+1, j)
和(nums1[i] + nums2[j+1], i, j+1)
分别加入堆中。注意要确保(i+1, j)
和(i, j+1)
没有超出数组范围。 - 将弹出的元素
(nums1[i], nums2[j])
加入结果数组。
- 弹出堆顶元素
代码
/**
* @param {number[]} nums1
* @param {number[]} nums2
* @param {number} k
* @return {number[][]}
*/
var kSmallestPairs = function (nums1, nums2, k) {
let heap = [];
const add = (val) => {
heap.push(val);
heapifyUp(heap.length - 1);
};
const pop = () => {
if (heap.length == 0) {
return null;
}
const top = heap[0];
const last = heap.pop();
if (heap.length > 0) {
heap[0] = last;
heapifyDown(0);
}
return top;
};
const heapifyUp = (i) => {
while (i) {
const parent = Math.floor((i - 1) / 2);
if (heap[i][0] < heap[parent][0]) {
[heap[i], heap[parent]] = [heap[parent], heap[i]];
i = parent;
} else {
break;
}
}
};
const heapifyDown = (i) => {
const left = i * 2 + 1;
const right = i * 2 + 2;
let min = i;
if (left < heap.length && heap[left][0] < heap[min][0]) {
min = left;
}
if (right < heap.length && heap[right][0] < heap[min][0]) {
min = right;
}
if (min !== i) {
[heap[i], heap[min]] = [heap[min], heap[i]];
heapifyDown(min);
}
};
let res = [];
if (!nums1 || !nums2 || nums1.length === 0 || nums2.length === 0 || k <= 0) {
return res;
}
for (let i = 0; i < Math.min(nums1.length, k); i++) {
add([nums1[i] + nums2[0], i, 0]);
}
while (k-- > 0 && heap.length > 0) {
const [sum, i, j] = pop();
res.push([nums1[i], nums2[j]]);
if (j + 1 < nums2.length) {
add([nums1[i] + nums2[j + 1], i, j + 1]);
}
}
return res;
};
相关题目
- [378. 有序矩阵中第 K 小的元素](https://leetcode.com/problems/kth-smallest-element-in-a-sorted-matrix)
- [719. 找出第 K 小的数对距离](https://leetcode.com/problems/find-k-th-smallest-pair-distance)
- [2040. 两个有序数组的第 K 小乘积](https://leetcode.com/problems/kth-smallest-product-of-two-sorted-arrays)