The original article is available at https://dawnmagnet.top/leetcode/2022/01/14/.
Given two integer arrays nums1 and nums2 sorted in ascending order, and an integer k, define a pair (u,v) whose first element comes from nums1 and whose second element comes from nums2.
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 in the sequence are:
[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 in the sequence are:
[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 pairs in the sequence may be returned: [1,3],[2,3]
Constraints:
1 <= nums1.length, nums2.length <= 10E5-10E9 <= nums1[i], nums2[i] <= 10E9nums1andnums2are both sorted in ascending order.1 <= k <= 1000
Source: LeetCode Link: https://leetcode-cn.com/problems/find-k-pairs-with-smallest-sums Copyright belongs to LeetCode. For commercial use, contact the official copyright holder; for non-commercial use, cite the source as required.
Analysis#
The main idea is to find the smallest k elements. Problems asking for the smallest or largest k elements commonly use a heap. A heap may also be called a min-heap, max-heap, or priority queue; fundamentally, it maintains the heap property. In a max-heap, every parent is greater than its children, while in a min-heap every parent is smaller than its children. While maintaining this property, insertion and removal take only \(O(lg(n))\) time, allowing us to obtain several extreme values efficiently.
Here, there are \(m*n\) possible values in total, because every pairwise sum of an element from the two arrays could be part of the result.
We maintain a max-heap called heap, whose size never exceeds k, the number of pairs required in the final answer. We first use two nested loops to add every possible pair to heap. Whenever the heap grows beyond k, remove its maximum. This still guarantees that the smallest k elements can be retained.
There are two useful optimizations (without them, the solution may time out):
- From the constraints,
nums1andnums2may each contain up to \(10^5\) elements, while the answer contains at most1000pairs. There is no need to inspect every combination, so limit both loop variables to at mostk. Every pair in the answer must come from pairing one of the firstkelements ofnums1with one of the firstkelements ofnums2. - During the traversal, we visit values from left to right. If the current value is greater than or equal to the maximum element in the heap, some combinations in the current inner loop no longer need to be examined, so we can
breakand skip them.
Rust Code#
use std::collections::BinaryHeap;
impl Solution {
pub fn k_smallest_pairs(nums1: Vec<i32>, nums2: Vec<i32>, k: i32) -> Vec<Vec<i32>> {
let mut res = vec![];
let mut heap = BinaryHeap::new();
for i in 0..(nums1.len()).min(k as usize) {
for j in 0..(nums2.len()).min(k as usize) {
heap.push((nums1[i] + nums2[j], nums1[i], nums2[j]));
if heap.len() > k as usize {
if let Some((a, b, c)) = heap.pop() {
if b == nums1[i] && c == nums2[j] {
break;
}
}
}
}
}
while heap.len() > 0 {
if let Some((a, b, c)) = heap.pop() {
res.push(vec![b, c]);
}
}
res
}
}
