查找和最小的 K 对数字

标签: 数组 堆（优先队列）

难度: Medium

给定两个以升序排列的整数数组 nums1 和 nums2 , 以及一个整数 k 。

定义一对值 (u,v)，其中第一个元素来自 nums1，第二个元素来自 nums2 。

请找到和最小的 k 个数对 (u1,v1),  (u2,v2)  ...  (uk,vk) 。

示例 1:

输入: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
输出: [1,2],[1,4],[1,6]
解释: 返回序列中的前 3 对数：
    [1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]

示例 2:

输入: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
输出: [1,1],[1,1]
解释: 返回序列中的前 2 对数：
     [1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]

示例 3:

输入: nums1 = [1,2], nums2 = [3], k = 3 
输出: [1,3],[2,3]
解释: 也可能序列中所有的数对都被返回:[1,3],[2,3]

提示:

  • 1 <= nums1.length, nums2.length <= 104
  • -109 <= nums1[i], nums2[i] <= 109
  • nums1, nums2 均为升序排列
  • 1 <= k <= 1000

注意：本题与主站 373 题相同：https://leetcode-cn.com/problems/find-k-pairs-with-smallest-sums/

Submission

运行时间: 32 ms

内存: 16.5 MB

class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        m, n = len(nums1), len(nums2)
        ans = []
        pq = [(nums1[i] + nums2[0], i, 0) for i in range(min(k, m))]
        while pq and len(ans) < k:
            _, i, j = heappop(pq)
            ans.append([nums1[i], nums2[j]])
            if j + 1 < n:
                heappush(pq, (nums1[i] + nums2[j + 1], i, j + 1))
        return ans

Explain

本题解采用了最小堆（优先队列）的方法来解决问题。首先，将 nums1 中的前 k 个元素与 nums2 中的第一个元素组成的数对，以及它们的和作为优先级，加入最小堆中。然后，每次从堆中取出和最小的数对，将其加入答案列表中。同时，如果该数对的 nums2 元素不是 nums2 的最后一个元素，则将该数对的 nums1 元素与 nums2 中下一个元素组成的新数对加入堆中。重复这个过程，直到找到 k 对数或堆为空。

时间复杂度: O(k log k)

空间复杂度: O(k)

class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        m, n = len(nums1), len(nums2)
        ans = []
        pq = [(nums1[i] + nums2[0], i, 0) for i in range(min(k, m))]
        while pq and len(ans) < k:
            _, i, j = heappop(pq)  # 从堆中取出和最小的数对
            ans.append([nums1[i], nums2[j]])  # 加入答案列表
            if j + 1 < n:
                heappush(pq, (nums1[i] + nums2[j + 1], i, j + 1))  # 将新的数对加入堆中
        return ans

Explore

这种策略是为了避免在堆初始化阶段就将所有可能的数对组合加入堆中，从而导致空间复杂度和时间复杂度过高。初始化时只考虑nums1的元素与nums2第一个元素的组合，是因为这样可以保证堆中始终维持最小的数对和，每次从堆中取出的都是当前已知的和最小的数对。之后每取出一个数对，再逐步考虑该数对的nums1元素与nums2中的下一个元素的组合，这样既控制了堆的大小，也能保证逐步找到k个最小和的数对。如果从一开始就考虑所有组合，堆的大小将难以控制，且初始化和处理的时间复杂度都会显著增加。

当从堆中取出一个数对后，考虑该数对的nums1元素与nums2中的下一个元素组成新的数对，是为了保证逐步扩展搜索空间，并且避免重复添加相同的数对。如果同时考虑nums1中的下一个元素，会导致重复和不必要的计算，因为对于每个nums1中的元素，它与nums2中每个元素的组合在初始化堆的时候已经按需逐步被加入了。这种方法确保了每次扩展都是基于当前已知的最小数对和，且每次扩展都有序且不重复，从而有效控制了堆的大小和算法的执行效率。

如果nums1或nums2是空数组，则没有有效的数对可以形成。在这种情况下，算法在开始执行前应该先检查nums1和nums2的长度。如果其中任何一个数组为空，直接返回一个空列表作为结果。这种预先检查可以避免进一步的无效操作，并且保证算法的鲁棒性。例如，在算法的开始添加如下代码：`if not nums1 or not nums2: return []`，这样可以在数据不满足条件时，及时返回结果，避免后续的无效计算。