给定两个以 升序排列 的整数数组 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,这样做的效率要高得多。