菜鸟的每日力扣系列——373. 查找和最小的 K 对数字

2022-12-12 17:38:11 浏览数 (1)

力扣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

0 人点赞