leetcode(力扣) 373. 查找和最小的 K 对数字(伪暴力法 & 优先队列(小根堆)---详细步骤解析)

题目描述

给定两个以 升序排列 的整数数组 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]

思路分析

法一

这道题是可以暴力的，不过不能纯暴力。
我们直接俩指针两两组合将全部可能都找出来然后排序，去前k个返回就行了。
报错了，超时了，看了一下输入的用例。

可以看到，输入的两个数组都是含有1W个数，但是k确是1000。
所以我们没必要吧全部的可能都找出来，这样会浪费很多时间。按照这个用例优化暴力算法，将两个数组都取前k个进行组合 可以得到 k^k 个数据 然后再排序输出前k个。
优化之后的暴力代码可以通过，时间复杂度大概超过4%。

完整代码:

class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        res = []
        for i in range(min(len(nums1),k)):
            for j in range(min(len(nums2),k)):
                res.append([nums1[i],nums2[j]])
        res = sorted(res,key = lambda x:(x[0]+x[1]))
        if len(res) < k:
            return res
        return res[0:k]

法二

如有如下数组 ：

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

显然，最小组合必是 nums1[0]+nums2[0]

我们设立一个中间数组 temp用于存放可能的数据。

• 将最小组合 nums1[0],nums2[0] 放入中间数组 temp。
  此时temp中仅有一个数据组合，所以直接加入到答案集。

• 然后找nums1[0],nums2[0]之后的下一个最小数，可能是nums1[0]+nums2[1] 或nums1[1]+nums2[0]。

• 将可能的这两组数据组合加入到中间数组temp。
  此时temp中的数据为 [ [2,3],[1,4] ] 这两个都是5，所以那个都可以，这里以 [2,3]在前。

• 继续展开[2,3], 按照前面的规律[2,3]之后的最小值可能是nums1[1]+nums2[1]与 nums1[2],nums2[0] 既[2,4]和[3,7]。将可能的数加入temp。此时的temp = 【【1，4】，【2，4】，【3，7】】

• 继续展开 【1，4】 按照上面的规律 【1，4】后面是【1，6】和【2，4】这里注意，【2，4】在上面已经被使用过了，所以这里不需要再重复计算，所以还要设立一个数组记录已经计算过的数组组合。

上面的过程中不难发现，我们的中间数组temp每次都将取出一个最小组合拿去展开。所以使用小根堆(优先队列)维护这个中间数组，让每次加入新数据之后，其根节点root保存的是组最小值，即temp[0]

不断重复上述过程即可，遵循 取1存2的步骤。

完整代码

class Solution:
    def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
        import heapq

        res = []
        # heap里存储 两者的和，nums1的下标，nums2的下标
        heap = [(nums1[0]+nums2[0],0,0)]
        repeat = set() # 记录重复组合
        while  len(res) < k:
        	if not heap:
                return res 
            summ,i,j = heapq.heappop(heap) # 取出最小值信息
            # sum为和， i 和j分别记录 nums1和nums2的下标值

            if (i,j) in repeat: # 重复值去掉
                continue
            repeat.add((i,j)) 
            res.append([nums1[i],nums2[j]])
            
			# 按照展开思路  取1存2
            if i+1<len(nums1):
                heapq.heappush(heap,(nums1[i+1]+nums2[j],i+1,j))
            if j +1 <len(nums2):
                heapq.heappush(heap,(nums1[i]+nums2[j+1],i,j+1))

        return res