力扣373. 查找和最小的 K 对数字
解题思路:多路归并的问题可以尝试用堆来解。对于本题可以先依次将nums1[0] nums2[0],nums1[1] nums2[0]……nums1[len(nums1)-1] nums2[0] push到堆中,由于数组是升序排列,其中nums1[0] nums2[0]一定是最小的,次小值是min(nums1[0] nums2[1], nums1[1] nums2[0])。
既然是每次去拿到最小值,我们可以使用小根堆来做,在python中使用heapq默认是小根堆。那么第一个入堆并从堆中弹出的答案是nums1[0] nums2[0],再让nums1[0] nums2[1]入堆,弹出第二个答案,以此类推;然后考虑取k对数字怎么实现,我们可以直接动态生成k个,那么循环条件应为当堆不为空且len(pairs)<7时动态入堆和出堆(pairs是存放最终答案的列表)。
上述的情况是确定了nums1去遍历nums2,接下来就是确定nums2来加nums1。但是我们发现,结果是加重复了,所以这里需要加上限制,当加到nums1[i 1] nums2[0]时,代表nums1的下一个加上nums2的第一个,为了避免重复加nums2的第一个,我们把nums[j]直接置为0。
代码语言:javascript复制from typing import List
from heapq import heappop, heappush
def k_smallest_pairs(nums1: List[int], nums2: List[int], k: int) -> List[int]:
h: list = []
def push(i, j):
if i < len(nums1) and j < len(nums2):
heappush(h, [nums1[i] nums2[j], i, j])
push(0, 0)
pairs: list = []
while h and len(pairs) < k:
_, i, j = heappop(h)
# print("i=", i, "j=", j)
pairs.append([nums1[i], nums2[j]])
push(i, j 1)
if j == 0: # 加回去[i 1],[0]对应nums1的下一个和nums2的第一个,避免重复加nums2的第一个
push(i 1, 0)
# print("pairs=", pairs)
return pairs
nums1 = [1,7,11]
nums2 = [2,4,6]
k = 3
print(k_smallest_pairs(nums1, nums2, k))
# i= 0 j= 0
# pairs= [[1, 2]]
# i= 0 j= 1
# pairs= [[1, 2], [1, 4]]
# i= 0 j= 2
# pairs= [[1, 2], [1, 4], [1, 6]]
END