Find K paris with smallest sums
用堆就行了,注意k可能大于所有组合数
class Solution(object):
def kSmallestPairs(self, nums1, nums2, k):
"""
:type nums1: List[int]
:type nums2: List[int]
:type k: int
:rtype: List[List[int]]
"""
n, m = len(nums1), len(nums2)
if not n or not m:
return []
result = []
checker = set()
heap = [(nums1[0] + nums2[0], 0, 0,)]
checker.add((0, 0,))
k = k if k <= n * m else n * m
while k:
value, i, j = heapq.heappop(heap)
result.append([nums1[i], nums2[j]])
if (i + 1, j,) not in checker and (i + 1 < n):
checker.add((i + 1, j,))
heapq.heappush(heap, (nums1[i + 1] + nums2[j], i + 1, j,))
if (i, j + 1,) not in checker and (j + 1 < m):
checker.add((i, j + 1,))
heapq.heappush(heap, (nums1[i] + nums2[j + 1], i, j + 1,))
k -= 1
return result