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

才浅Coding攻略

发布于 2022-12-12 17:38:11

4720

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

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

本文参与腾讯云自媒体同步曝光计划，分享自微信公众号。

原始发表：2022-01-20，如有侵权请联系 cloudcommunity@tencent.com 删除

python