Find k pairs with smallest sums

December 5, 2022

heap

Problem URL: Find k pairs with smallest sums

We will use a min heap to store the smallest sum. We will add the first element of each array to the heap. Then we will pop the smallest element from the heap and add the next element of the array from which the smallest element was popped. We will repeat this process until we have k pairs.

class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        if not nums1 or not nums2:
            return []
        heap = []
        for i in range(len(nums1)):
            heapq.heappush(heap, (nums1[i] + nums2[0], i, 0))
        res = []
        while heap and len(res) < k:
            _, i, j = heapq.heappop(heap)
            res.append([nums1[i], nums2[j]])
            if j + 1 < len(nums2):
                heapq.heappush(heap, (nums1[i] + nums2[j + 1], i, j + 1))
        return res

Time complexity: O(klog(n))
Space complexity: O(n)