Find K Pairs with Smallest Sums
Intuition
Solution
class Solution:
@staticmethod
def _add_to_combine(combine, pair, value):
if value not in combine:
combine[value] = []
combine[value].append(pair)
@staticmethod
def _remove_from_combine(combine, value):
pair = combine[value].pop()
if not combine[value]:
combine.pop(value)
return pair
def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
if not nums1:
return []
combine = {} # Key is the sum and Value is a list of possibles to sum to it
heap = []
for j in range(len(nums2)):
heapq.heappush(heap, nums1[0] + nums2[j])
self._add_to_combine(combine, (0, j), nums1[0] + nums2[j])
ans = []
for _ in range(k):
# print(heap, combine)
if not combine:
break
value = heapq.heappop(heap)
i, j = self._remove_from_combine(combine, value)
ans.append([nums1[i], nums2[j]])
i += 1
if i < len(nums1):
heapq.heappush(heap, nums1[i] + nums2[j])
self._add_to_combine(combine, (i, j), nums1[i] + nums2[j])
return ans