题目描述
给定2个以升序排列的整数数组和1个整数k,从2个数组中各拿出一个数字组成一对,找出k对和最小的数字组合。
解法1:最简单的想法就是暴力brute force解法,但效率肯定不高。
样例
Input: nums1 = [1,1,2], nums2 = [1,2,3], k = 2
Output: [1,1],[1,1]
Explanation: The first 2 pairs are returned from the sequence:
[1,1],[1,1],[1,2],[2,1],[1,2],[2,2],[1,3],[1,3],[2,3]
算法1
heap
解法2: 最小堆。
把所有的点对加入到最小堆,然后输出前k个。但没有利用到“两个数组都有序”这个条件,就算数组无序,也可以利用这个方法。要利用有序这个条件,可以借助mergesort的思路,pair的第一个元素至多包含了nums1数组的前k个元素,k以后的可以不用考虑。所以,这形成了k个list,每一个list都包含了nums2的元素。每一次取所有list中的最小值,然后该list下一个元素入队。
时间复杂度
python3 代码
from heapq import heappush, heappop
class Solution(object):
def kSmallestPairs(self, nums1, nums2, k):
"""
:type nums1: List[int]
:type nums2: List[int]
:type k: int
:rtype: List[List[int]]
"""
pairs = []
if len(nums1) > len(nums2):
tmp = self.kSmallestPairs(nums2, nums1, k)
for pair in tmp:
pairs.append([pair[1], pair[0]])
return pairs
min_heap = []
def push(i, j):
if i < len(nums1) and j < len(nums2):
heappush(min_heap, [nums1[i] + nums2[j], i, j])
push(0, 0)
while min_heap and len(pairs) < k:
_, i, j = heappop(min_heap) # j 由上一次pop出来的数字产生
pairs.append([nums1[i], nums2[j]])
push(i, j + 1) # j=j+1, 只会越来越大,不会浪费时间
if j == 0:
push(i + 1, 0) # at most queue min(n, m) space
return pairs