03_FindKPairsWithSmallestSum.py

# Question link - https://leetcode.com/problems/find-k-pairs-with-smallest-sums/description/?envType=study-plan-v2&envId=top-interview-150

class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        # For this problem , we are going to use heap and visited 
        res = [] # for the result 

        # Base case 
        if not nums1 or not nums2 or not k:
            return res

        heap = []
        visited = set()

        heapq.heappush(heap , (nums1[0]+nums2[0],0,0))
        visited.add((0 , 0))

        while k and heap:
            _ ,i,j = heapq.heappop(heap)
            res.append([nums1[i] , nums2[j]])

            if i+1 < len(nums1) and (i+1 , j) not in visited:
                heapq.heappush(heap , (nums1[i+1] + nums2[j] , i+1 , j))
                visited.add((i+1 , j))

            if j+1 < len(nums2) and (i , j+1) not in visited:
                heapq.heappush(heap , (nums1[i] + nums2[j+1], i , j+1))
                visited.add((i , j+1))

            k -= 1
        
        return res