题目:
题解:
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