题目
给定两个以 非递减顺序排列 的整数数组
nums1
和nums2
, 以及一个整数k
。定义一对值
(u,v)
,其中第一个元素来自nums1
,第二个元素来自nums2
。请找到和最小的
k
个数对(u1,v1)
,(u2,v2)
...(uk,vk)
。
解题
python
from heapq import heappush, heappop
def k_smallest_pairs(nums1, nums2, k):
# 存储结果的列表
result = []
# 边界条件
if not nums1 or not nums2:
return result
# 使用最小堆
min_heap = []
# 初始化最小堆,存储 (和, nums1中的索引, nums2中的索引)
for i in range(min(len(nums1), k)): # 只需要前 k 个元素进行初始化
heappush(min_heap, (nums1[i] + nums2[0], i, 0))
while k > 0 and min_heap:
current_sum, i, j = heappop(min_heap)
result.append((nums1[i], nums2[j]))
# 如果还有下一个元素,则将下一个元素的对加到堆中
if j + 1 < len(nums2):
next_sum = nums1[i] + nums2[j + 1]
heappush(min_heap, (next_sum, i, j + 1))
k -= 1
return result
# 测试代码
nums1 = [1, 7]
nums2 = [3, 4]
k = 3
print(f"和最小的{k}个数对 : {k_smallest_pairs(nums1, nums2, k)}")
和最小的3个数对 : [(1, 3), (1, 4), (7, 3)]