You are given two integer arrays nums1 and nums2 sorted in ascending 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 <= 105
- -109 <= nums1[i], nums2[i] <= 109
- nums1 and nums2 both are sorted in ascending order.
- 1 <= k <= 104
又是这种合并排序的问题， 我们还是需要建一个池子保存当前备选的选项， 这些选项由三个要素组成， index1 为 nums1 的 index， index2 为 nums2 的 index， sum 为 nums1[index1] + nums2[index2], 这三个当中 sum 是为排序准备的， 我们遍历 nums1 生成(nums1[i]+nums2[0], i, 0)放入一个 binary heap 中， 这样我们就得到了所有 nums1 中的元素与 nums2[0]所组合出来的备选项， 我们每次从中拿出最小的那个， 将 nums1[index1]和 nums2[index2]放到答案中， 然后我们将 index2+1 并重新计算备选项， 把备选项重新放到 binary heap 中就可以了。要注意的是 k 可能大于所有备选项的数量
use std::cmp::Reverse;
use std::collections::BinaryHeap;
impl Solution {
pub fn k_smallest_pairs(nums1: Vec<i32>, nums2: Vec<i32>, mut k: i32) -> Vec<Vec<i32>> {
let mut heap = nums1
.iter()
.enumerate()
.map(|(i, &n)| Reverse((n + nums2[0], i, 0usize)))
.collect::<BinaryHeap<Reverse<(i32, usize, usize)>>>();
let mut ans = Vec::new();
while k > 0 {
if let Some(Reverse((_, i1, mut i2))) = heap.pop() {
ans.push(vec![nums1[i1], nums2[i2]]);
k -= 1;
i2 += 1;
if i2 == nums2.len() {
continue;
}
heap.push(Reverse((nums1[i1] + nums2[i2], i1, i2)));
continue;
}
break;
}
ans
}
}