373. Find K Pairs with Smallest Sums

• categories: Code, Interview Question, leetcode, Medium
• source: https://leetcode.com/problems/find-k-pairs-with-smallest-sums
• topics: Heap, Array

Description

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

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

Idea

• Heap and keep visited in the set
• Avoid duplicates without visited set. If the chosen pair is the first one in its row, then the first pair in the next row is added to the queue

# # # # # ? . .
# # # ? . . . .
# ? . . . . . .   "#" means pair already in the output
# ? . . . . . .   "?" means pair currently in the queue
# ? . . . . . .
? . . . . . . .
. . . . . . . .

Code

Heap and set

class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        ans, visited = [], set()
        heap = [(nums1[0] + nums2[0], 0, 0)]
        while len(ans) < k:
            _, i, j = heappop(heap)
            ans.append([nums1[i], nums2[j]])
            
            if i + 1 < len(nums1) and (i + 1, j) not in visited:
                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:
                heappush(heap, (nums1[i] + nums2[j + 1], i, j + 1))
                visited.add((i, j + 1))
 
        return ans

Heap without set

class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        ans = []
        heap = [(nums1[0] + nums2[0], 0, 0)]
        while len(ans) < k:
            _, i, j = heappop(heap)
            ans.append([nums1[i], nums2[j]])
            
            if (j + 1) < len(nums2):
                heappush(heap, (nums1[i] + nums2[j+1], i, j+1))
 
            if j == 0 and (i+1) < len(nums1):
                heappush(heap, (nums1[i+1] + nums2[0], i+1, 0))
 
        return ans