文章目录
- 题目描述
- 思路分析
- 法一
- 完整代码:
- 法二
- 完整代码
题目描述
给定两个以 升序排列 的整数数组 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个进行组合 可以得到 kk 个数据 然后再排序输出前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:
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
