给定两个以 升序排列 的整数数组 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]
1. 优先级队列
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
2. 二分法
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
