给定两个以 升序排列 的整数数组 nums1 和 nums2 , 以及一个整数 k 。
定义一对值 (u,v),其中第一个元素来自 nums1,第二个元素来自 nums2 。
请找到和最小的 k 个数对 (u1,v1), (u2,v2) … (uk,vk)。
官方解法:
方法一:优先队列
已知含有 n 个数对的时候,那么第 n+1 个数对必定从以下组合中产生,
(a1+1,b1), (a1,b1+1), (a2+1,b2), (a2,b2+1), ... , (an,bn+1)
其中 a1,a2 代表的是在升序 nums1, nums2 的位序。
算法实现方面使用
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
方法二:二分查找
class Solution:
def kSmallestPairs(self, nums1: List[int], nums2: List[int], k: int) -> List[List[int]]:
m, n = len(nums1), len(nums2)
# 二分查找第 k 小的数对和
left, right = nums1[0] + nums2[0], nums1[m - 1] + nums2[n - 1] + 1
while left < right:
mid = (left + right) // 2
cnt = 0
i, j = 0, n - 1
while i < m and j >= 0:
if nums1[i] + nums2[j] > mid:
j -= 1
else:
cnt += j + 1
i += 1
if cnt < k:
left = mid + 1
else:
right = mid
pairSum = left
ans = []
# 找数对和小于 pairSum 的数对
i = n - 1
for num1 in nums1:
while i >= 0 and num1 + nums2[i] >= pairSum:
i -= 1
for j in range(i + 1):
ans.append([num1, nums2[j]])
if len(ans) == k:
return ans
# 找数对和等于 pairSum 的数对
i = n - 1
for num1 in nums1:
while i >= 0 and num1 + nums2[i] > pairSum:
i -= 1
j = i
while j >= 0 and num1 + nums2[j] == pairSum:
ans.append([num1, nums2[j]])
if len(ans) == k:
return ans
j -= 1
return ans
来源:力扣(LeetCode)
链接:https://leetcode-cn.com/problems/find-k-pairs-with-smallest-sums
著作权归领扣网络所有。商业转载请联系官方授权,非商业转载请注明出处。
Python小知识
堆
实际上,Python没有独立的堆类型,而只有一个包含一些堆操作函数的模块。这个模块名为heapq(其中的q表示队列),它包含6个函数,其中前4个与堆操作直接相关。必须使用列表来表示堆对象本身。
优先队列:import heapq 堆的使用方式:heapq.heappop() 是压出最小的元素,heapq.heappush()是压进堆。
元素的排列顺序并不像看起来那么随意。它们虽然不是严格排序的,但必须保证一点:位置i处的元素总是大于位置i // 2处的元素(反过来说就是小于位置2 * i和2 * i + 1处的元素)。这是底层堆算法的基础,称为堆特征(heap property)。
函数heappop弹出最小的元素(总是位于索引0处),并确保剩余元素中最小的那个位于索引0处(保持堆特征)。虽然弹出列表中第一个元素的效率通常不是很高,但这不是问题,因为heappop会在幕后做些巧妙的移位操作。
优先队列让你能够以任意顺序添加对象,并随时(可能是在两次添加对象之间)找出(并删除)最小的元素。相比于列表方法min,这样做的效率要高得多。