You are given two integer arrays nums1 and nums2 sorted in ascending 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.

Find the k pairs (u1,v1),(u2,v2) ...(uk,vk) with the smallest sums.


Example 1:

Given nums1 = [1,7,11], nums2 = [2,4,6],  k = 3

Return: [1,2],[1,4],[1,6]

The first 3 pairs are returned from the sequence:
[1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]



Example 2:

Given nums1 = [1,1,2], nums2 = [1,2,3],  k = 2

Return: [1,1],[1,1]

The first 2 pairs are returned from the sequence:
[1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]



Example 3:

Given nums1 = [1,2], nums2 = [3],  k = 3 

Return: [1,3],[2,3]

All possible pairs are returned from the sequence:
[1,3],[2,3]



Credits:Special thanks to @elmirap and @StefanPochmann for adding this problem and creating all test cases.

Difficulty: interview

Starter code:

class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        

Solutions:
from heapq import heappush, heappop
 
 class Solution:
     def kSmallestPairs(self, nums1, nums2, k):
         """
         :type nums1: List[int]
         :type nums2: List[int]
         :type k: int
         :rtype: List[List[int]]
         """
         heap = []
         def push(i, j):
             if i < len(nums1) and j < len(nums2):
                 heappush(heap, (nums1[i] + nums2[j], i, j))
         push(0,0)
         res = []
         while heap and len(res) < k:
             _, i, j = heappop(heap)
             res.append([nums1[i], nums2[j]])
             push(i, j+1)
             if j == 0:
                 push(i+1, 0)
         return res

class Solution:
     def kSmallestPairs(self, nums1, nums2, k):
         """
         :type nums1: List[int]
         :type nums2: List[int]
         :type k: int
         :rtype: List[List[int]]
         """
         queue = []
         def push(i, j):
             if i < len(nums1) and j < len(nums2):
                 heapq.heappush(queue, [nums1[i] + nums2[j], i, j])
         push(0, 0)
         pairs = []
         while queue and len(pairs) < k:
             _, i, j = heapq.heappop(queue)
             pairs.append([nums1[i], nums2[j]])
             push(i, j + 1)
             if j == 0:
                 push(i + 1, 0)
         return pairs

class Solution:
     def kSmallestPairs(self, nums1, nums2, k):
         """
         :type nums1: List[int]
         :type nums2: List[int]
         :type k: int
         :rtype: List[List[int]]
         """
         if not nums1 or not nums2 or k < 1:
             return []
         n1, n2 = len(nums1), len(nums2)
         stack = [(nums1[0] + nums2[0], (0, 0))]
         visited = set()
         res = []
         while len(res) < k and stack:
             num, (p1, p2) = heapq.heappop(stack)
             res.append([nums1[p1], nums2[p2]])
             if p1 + 1 < n1 and (p1 + 1, p2) not in visited:
                 heapq.heappush(stack, (nums1[p1 + 1] + nums2[p2], (p1 + 1, p2)))
                 visited.add((p1 + 1, p2))            
             if p2 + 1 < n2 and (p1, p2 + 1) not in visited:
                 heapq.heappush(stack, (nums1[p1] + nums2[p2 + 1], (p1, p2 + 1)))
                 visited.add((p1, p2 + 1))            
         return res

class Solution:
     def kSmallestPairs(self, nums1, nums2, k):
         """
         :type nums1: List[int]
         :type nums2: List[int]
         :type k: int
         :rtype: List[List[int]]
         """
         #heap based solution
         #visualize as a matrix
         #best possible solution will be first row or first column
         #O(2K + Klog2K) time; O(1) time
         
         if not nums1 or not nums2:
             return []
         
         queue = [(nums1[0]+nums2[0],0,0)]
         ans = []
         while queue and len(ans) < k:
             _, i, j = heapq.heappop(queue)
             ans.append([nums1[i],nums2[j]])            
             if j+1 < len(nums2): #smallest values #row-wise
                 heapq.heappush(queue, (nums1[i]+nums2[j+1],i,j+1))                
             if j == 0 and i+1 < len(nums1): #smallest values #column-wise #takes care of row traversal                     
                 heapq.heappush(queue, (nums1[i+1]+nums2[j],i+1,j))                                           
         return ans
             


import heapq
 class Solution:
     def kSmallestPairs(self, nums1, nums2, k):
         """
         :type nums1: List[int]
         :type nums2: List[int]
         :type k: int
         :rtype: List[List[int]]
         """
         pq = []
         m, n = len(nums1), len(nums2)   # get lenth
         if m == 0 or n == 0:            # if one of the list is empty return []
             return []
         for i in range(m):              # get through the first list
             if len(pq) == k and pq[0][0] > -nums1[i]-nums2[0]: # stop if larger than current k
                 break
             for j in range(n):          # go throught the second list
                 s = -nums1[i]-nums2[j]   # get the sum
                 if len(pq) < k or pq[0][0] < s: # if less than k or less than current k
                     heapq.heappush(pq, (s, [nums1[i], nums2[j]]))   # push 
                     if len(pq) > k: # get it out
                         heapq.heappop(pq)
                 else:                   # stop if larger than current k
                     break
         return [item[1] for item in pq]
         

class Solution(object):
     def kSmallestPairs(self, nums1, nums2, k, heap=[]):
         for n1 in nums1:
             for n2 in nums2:
                 if len(heap) < k: heapq.heappush(heap, (-n1-n2, [n1, n2]))
                 else:
                     if heap and -heap[0][0] > n1 + n2:
                         heapq.heappop(heap)
                         heapq.heappush(heap, (-n1-n2, [n1, n2]))
                     else: break
         return [heapq.heappop(heap)[1] for _ in range(k) if heap]

import heapq
 
 class Solution:
     def kSmallestPairs(self, nums1, nums2, k):
         """
         :type nums1: List[int]
         :type nums2: List[int]
         :type k: int
         :rtype: List[List[int]]
         """
         
         if not nums1 or not nums2:
             return []
         
         index = [0 for _ in range(len(nums1))]
         q = [[nums1[i] + nums2[index[i]], i] for i in range(len(nums1))]
         heapq.heapify(q)
         i = 0
         ans = []
         while i < k and i < len(nums1) * len(nums2):
             pair = heapq.heappop(q)
             ans.append([nums1[pair[1]], nums2[index[pair[1]]]])
             i += 1
             if index[pair[1]] + 1 < len(nums2):
                 index[pair[1]] += 1
                 heapq.heappush(q, [nums1[pair[1]] + nums2[index[pair[1]]], pair[1]])
             
         return ans
             
         

class Solution:
     def kSmallestPairs(self, nums1, nums2, k):
         """
         :type nums1: List[int]
         :type nums2: List[int]
         :type k: int
         :rtype: List[List[int]]
         """
         q = []
         def push(i,j):
             if i<len(nums1) and j<len(nums2):
                 heapq.heappush(q,[nums1[i]+nums2[j],i,j])
         push(0,0) # pushing first sum into heap
         res = []
         while q and len(res) <k:
             s,i,j = heapq.heappop(q) #pops out min till here
             res.append([nums1[i],nums2[j]])
             push(i,j+1)
             if j == 0:
                 push(i+1,0)
         return res        

class Solution:
     def kSmallestPairs(self, nums1, nums2, k):
         """
         :type nums1: List[int]
         :type nums2: List[int]
         :type k: int
         :rtype: List[List[int]]
         """
         result = []
         import heapq
         m = len(nums1)
         n = len(nums2)
         if m*n>0:
             used = [[0 for _ in range(n)] for _ in range(m)]
             used[0][0] = 1
             h = [(nums1[0]+nums2[0],(0,0))]
             while k  > 0 and h:
                 _,(s1,s2) = heapq.heappop(h)
                 result.append([nums1[s1],nums2[s2]])
                 if s1+1 < m and used[s1+1][s2] == 0:
                     heapq.heappush(h,(nums1[s1+1]+nums2[s2],(s1+1,s2)))
                     used[s1+1][s2] = 1
                 if s2+1 < n and used[s1][s2+1] == 0:
                     heapq.heappush(h,(nums1[s1]+nums2[s2+1],(s1,s2+1)))
                     used[s1][s2+1] = 1
                 k -= 1
         return result


class Solution:
     def kSmallestPairs(self, nums1, nums2, k):
         if not nums1 or not nums2:
             return []
         queue = []
         heapq.heappush(queue, [nums1[0] + nums2[0], 0, 0])
         ans = []
         seen = {(0, 0):1}
         while k > len(ans) and queue:
             top = heapq.heappop(queue)
             i, j = top[1], top[2]
             ans.append([nums1[i],nums2[j]])
             #print(i, j, queue)
             if i+1 < len(nums1) and (i+1, j) not in seen:
                 heapq.heappush(queue, [nums1[i+1] + nums2[j], i+1, j])
                 seen[(i+1, j)] = 1
             if j+1 < len(nums2) and (i, j+1) not in seen:
                 heapq.heappush(queue, [nums1[i] + nums2[j+1], i, j+1])
                 seen[(i, j+1)] = 1
         return ans