Find K Pairs with Smallest Sums

Problem: Find K Pairs with Smallest Sums

We can use a heap to store all pairs with sums as the keys. But in this way memory limit will get exceeded easily. To use less memory, we can first store the first number from first array pairing with all other numbers from the second array. Every time we pop one pair from the heap, we replace the number from first array with its next number and push the pair of numbers back.

Code in Python:

from heapq import *

class Solution(object):
    def kSmallestPairs(self, nums1, nums2, k):
        """
        :type nums1: List[int]
        :type nums2: List[int]
        :type k: int
        :rtype: List[List[int]]
        """
        n, m = len(nums1), len(nums2)
        if n is 0 or m is 0:
            return []
        heap, ksmallest = [(nums1[0] + nums2[i], 0, i) for i in range(m)], []
        heapify(heap)
        while len(ksmallest) < k and len(heap) > 0:
            _, i, j = heappop(heap)
            ksmallest.append([nums1[i], nums2[j]])
            if i + 1 < n:
                heappush(heap, (nums1[i+1] + nums2[j], i+1, j))
        return ksmallest