从两个升序排列的数组分别抽出一个数组成pair,输出和最小的k个pair,优先队列应用
1 #373
2 #Runtime: 1051 ms (Beats 45.18%)
3 #Memory: 29.9 MB (Beats 87.95%)
4
5 class Solution(object):
6 def kSmallestPairs(self, nums1, nums2, k):
7 """
8 :type nums1: List[int]
9 :type nums2: List[int]
10 :type k: int
11 :rtype: List[List[int]]
12 """
13 fg = set()
14 ans = []
15 hp = []
16 for i in range(min(len(nums1), k)):
17 heapq.heappush(hp, (nums1[i] + nums2[0], nums1[i], nums2[0], 0))
18 while k and hp:
19 _, i, j, idx = heapq.heappop(hp)
20 ans.append([i, j])
21 if idx < len(nums2) - 1:
22 heapq.heappush(hp, (i + nums2[idx + 1], i, nums2[idx + 1], idx + 1))
23 k -= 1
24 return ans