数组中最小的对和

题目描述

给定两个有序整数数组 nums1 和 nums2，请找出数组中最小的 k 个数对（u,v），其中第一个数组来自 nums1，第二个数组来自 nums2。

例如，给定如下两个数组：

nums1 = [1,7,11], nums2 = [2,4,6], k = 3

返回如下 3 个数对：

[1,2],[1,4],[1,6]

解题思路

方法1：暴力

可以将数对全部找出来并排序，然后取前k个。

时间复杂度：$O(n^2 \log n)$。

方法2：最小堆

在数组 nums1 中选一个数，再在数组 nums2 中选一个数，组成一个数对。这样得到的数字对的总数为 m*n，其中 m 和 n 为两个数组的长度。

现在的问题是，如何在所有的数对中找到最小的 k 个数对？一个显然的思路是将所有的数对排序，并取其前 k 个。但是，如果数组长度较大时，这样的时间复杂度会很高（$O(mn \log mn)$）。

为了优化时间复杂度，我们可以考虑一次取出多个数对，然后从中选出最小的数对。

具体地，我们可以使用一个最小堆来存储当前所有的包含 nums1 中的数和 nums2 中的数的数对。初始时，只将 (0,0) 加入到最小堆中。每次从最小堆中取出一个数对 (i,j)，并将其加入结果中，同时将 (i+1,j) 和 (i,j+1) 这两个数对加入最小堆中。如此重复 k 次，则可以得到前 k 个最小的数对。

此方法时间复杂度为 $O(k \log k)$。

代码示例

方法1：暴力

def kSmallestPairs(nums1, nums2, k):
    pairs = []
    for num1 in nums1:
        for num2 in nums2:
            pairs.append((num1, num2))
    pairs.sort(key=lambda x: x[0] + x[1])
    return pairs[:k]

方法2：最小堆

import heapq

def kSmallestPairs(nums1, nums2, k):
    if not nums1 or not nums2:
        return []
    pqueue = [(nums1[0]+nums2[0], 0, 0)]
    res = []
    while pqueue and len(res) < k:
        _, i, j = heapq.heappop(pqueue)
        res.append([nums1[i], nums2[j]])
        if i+1 < len(nums1) and j == 0:
            heapq.heappush(pqueue, (nums1[i+1]+nums2[0], i+1, 0))
        if j+1 < len(nums2):
            heapq.heappush(pqueue, (nums1[i]+nums2[j+1], i, j+1))
    return res

总结

本题解介绍了两种解法，暴力和最小堆。其中最小堆方法的时间复杂度为 $O(k \log k)$，优于暴力方法的时间复杂度 $O(mn \log mn)$，但需要额外的空间来存储堆。