373. 查找和最小的 K 对数字

exiaohu 于 2022-01-14 发布

题目链接:373. 查找和最小的 K 对数字

如题意,使用小根堆找到最小的数对即可。

需要注意的是,如果对两个数组直接做笛卡儿积,会超时,所以先分别找到两个数组中最小的 k 个数,然后对这两个较短数组做笛卡儿积。

from heapq import nsmallest
from typing import List
from itertools import product


class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        return nsmallest(k, ((u, v) for u, v in product(nsmallest(k, nums1), nsmallest(k, nums2))), key=sum)