↓ Skip to main content
LeetCode 373: Find K Pairs with Smallest Sums

LeetCode 373: Find K Pairs with Smallest Sums

·555 words·3 mins· loading · loading ·
Author
Dawn Magnet
Tech enthusiast & daily solver.

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] <= 10E9
  • nums1 and nums2 are 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, nums1 and nums2 may each contain up to \(10^5\) elements, while the answer contains at most 1000 pairs. There is no need to inspect every combination, so limit both loop variables to at most k. Every pair in the answer must come from pairing one of the first k elements of nums1 with one of the first k elements of nums2.
  • 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 break and 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
    }
}