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.
Find 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]
struct Solution;
use std::cmp::Reverse;
use std::collections::BinaryHeap;
use std::collections::HashSet;
impl Solution {
fn k_smallest_pairs(nums1: Vec<i32>, nums2: Vec<i32>, mut k: i32) -> Vec<Vec<i32>> {
let n1 = nums1.len();
let n2 = nums2.len();
let mut visited: HashSet<(usize, usize)> = HashSet::new();
let mut queue: BinaryHeap<(Reverse<i32>, usize, usize)> = BinaryHeap::new();
if 0 < n1 && 0 < n2 && visited.insert((0, 0)) {
queue.push((Reverse(nums1[0] + nums2[0]), 0, 0));
} else {
return vec![];
}
let mut res = vec![];
while k > 0 {
if let Some((_, i, j)) = queue.pop() {
res.push(vec![nums1[i], nums2[j]]);
if i + 1 < n1 && visited.insert((i + 1, j)) {
queue.push((Reverse(nums1[i + 1] + nums2[j]), i + 1, j));
}
if j + 1 < n2 && visited.insert((i, j + 1)) {
queue.push((Reverse(nums1[i] + nums2[j + 1]), i, j + 1));
}
k -= 1;
} else {
break;
}
}
res
}
}
#[test]
fn test() {
let nums1 = vec![1, 7, 11];
let nums2 = vec![2, 4, 6];
let k = 3;
let mut res = vec_vec_i32![[1, 2], [1, 4], [1, 6]];
let mut ans = Solution::k_smallest_pairs(nums1, nums2, k);
res.sort();
ans.sort();
assert_eq!(ans, res);
let nums1 = vec![1, 1, 2];
let nums2 = vec![1, 2, 3];
let k = 2;
let mut res = vec_vec_i32![[1, 1], [1, 1]];
let mut ans = Solution::k_smallest_pairs(nums1, nums2, k);
res.sort();
ans.sort();
assert_eq!(ans, res);
let nums1 = vec![1, 2];
let nums2 = vec![3];
let k = 3;
let mut res = vec_vec_i32![[1, 3], [2, 3]];
let mut ans = Solution::k_smallest_pairs(nums1, nums2, k);
res.sort();
ans.sort();
assert_eq!(ans, res);
}