給定兩個(gè)以 升序排列 的整數(shù)數(shù)組 nums1 和 nums2 , 以及一個(gè)整數(shù) k 胶哲。
定義一對(duì)值 (u,v)症歇，其中第一個(gè)元素來(lái)自 nums1魏铅，第二個(gè)元素來(lái)自 nums2 淮腾。
請(qǐng)找到和最小的 k 個(gè)數(shù)對(duì) (u1,v1), (u2,v2) ... (uk,vk) 糟需。
示例 1:
輸入: nums1 = [1,7,11], nums2 = [2,4,6], k = 3
輸出: [1,2],[1,4],[1,6]
解釋: 返回序列中的前 3 對(duì)數(shù)：
[1,2],[1,4],[1,6],[7,2],[7,4],[11,2],[7,6],[11,4],[11,6]
1. 優(yōu)先級(jí)隊(duì)列
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 小的數(shù)對(duì)和
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 = []
# 找數(shù)對(duì)和小于 pairSum 的數(shù)對(duì)
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
# 找數(shù)對(duì)和等于 pairSum 的數(shù)對(duì)
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
