373. Find K Pairs with Smallest Sums

Problem Statement

link: LeetCode.cn LeetCode 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:

Input: nums1 = [1,7,11], nums2 = [2,4,6], k = 3 Output: [[1,2],[1,4],[1,6]]

Input: nums1 = [1,1,2], nums2 = [1,2,3], k = 2 Output: [[1,1],[1,1]]

Solution Approach

Use a min-heap to efficiently track and retrieve the k pairs with the smallest sums by pushing initial pairs and iteratively adding the next possible pairs from the arrays.

Algorithm

1. Initialize Min-Heap: Create a min-heap to store pairs along with their sums, starting with the first element of nums2 paired with each element of nums1.
2. Iterate and Extract Minimum: Repeatedly extract the smallest sum pair from the heap, adding it to the result list, and then push the next possible pair involving the next element from nums2.
3. Limit Result Size: Continue the process until the result list contains k pairs or there are no more pairs to process.

Implement

    class Solution: def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
    m, n = len(nums1), len(nums2)
    res = []
    pq = [(nums1[i] + nums2[0]) for i in range(min(k, m))]
    while pq and len(res) < k:
        min_sum, i, j = heappop(pq)
        res.append([nums1[i], nums2[j]])
        if j + 1 < n:
            heappush(pq, (nums1[i] + nums2[j + 1], i, j + 1))
    return res   

    
    

Twitter  facebook  Email