LC 0373 [M] Find K Pairs with Smallest Sums - ALawliet/algorithms GitHub Wiki

Basic idea: Use min_heap to keep track on next minimum pair sum, and we only need to maintain K possible candidates in the data structure.

Some observations: For every numbers in nums1, its best partner(yields min sum) always strats from nums2[0] since arrays are all sorted; And for a specific number in nums1, its next candidate sould be [this specific number] + nums2[current_associated_index + 1], unless out of boundary;)

Here is a simple example demonstrate how this algorithm works.

The run time complexity is O(kLogk) since que.size <= k and we do at most k loop.

We create a priority queue using the sum as the priority. Note that python uses a min heap so the next item you pop off of the heap will be the coords with the smallest sum. We know that the smallest pair are the items at (0,0). We also know that the next smallest pair must be at either (1,0) or (0,1) since the inputs are sorted. Using induction you could prove that this holds for all any arbitrary coordinate, which is left to the reader.

When we push these items onto the heap, in logarithmic time the heap will put the smallest sum at the beginning of the heap. Therefore, whenever we pop something off of the heap we are guaranteed that it is the next smallest item.

We iterate until there are no more potential pairs or we have len(ret) == k.

Given the matrix for [1,7,11] and [2,4,6], we can do BFS from the top-left position to expand candidates from only right cell and bottom cell. Since we treat it as a graph, we need to skip visited vertices as I used a dictionary visited.

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
        
        min_heap = []
        visited = set()
        
        heappush(min_heap, (nums1[0] + nums2[0], (0, 0)))
        
        visited.add((0, 0))
        
        while len(res) < k and min_heap:
            _, ij = heappop(min_heap)
            i, j = ij
            res.append([nums1[i], nums2[j]])
            
            if i+1 < len(nums1) and (i+1, j) not in visited:
                heappush(min_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:
                heappush(min_heap, (nums1[i] + nums2[j+1], (i, j+1)))
                visited.add((i, j+1))
        return res

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
        
        min_heap = []
        visited = set()
        
        heappush(min_heap, (nums1[0] + nums2[0], 0, 0))
        
        visited.add((0, 0))
        
        while len(res) < k and min_heap:
            _, i, j = heappop(min_heap)
            res.append([nums1[i], nums2[j]])
            
            if i+1 < len(nums1) and (i+1, j) not in visited:
                heappush(min_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:
                heappush(min_heap, (nums1[i] + nums2[j+1], i, j+1))
                visited.add((i, j+1))
        return res