373. Find K Pairs with Smallest Sums

373. Find K Pairs with Smallest Sums

 

https://youtu.be/Youk8DDnLU8?si=dTtsC6RqlaGW-wWl

 

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

        if not nums1 or not nums2 or not k:
            return res

        heap = []
        visited = set()

        # 1번째 pair 추가: 정렬되어 있으므로, 둘다 첫번째, idx=0 이면, 합이 최소
        # 힙에 (sum, i, j) 추가되게.
        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]])

            # (i+1, j) 와 (i, j+1) 순서쌍을 힙에 추가해서 
            # 다음턴에서 더 최소값이 heappop 으로 튀어나오게.
            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