import heapq
def min_k_pair_sum(array1, array2, k):
if not array1 or not array2 or k == 0:
return 0
heap = []
visited = set()
# 初始时,最小的组合是array1[0] + array2[0]
heapq.heappush(heap, (array1[0] + array2[0], 0, 0))
visited.add((0, 0))
result = 0
for _ in range(k):
if not heap:
break
current_sum, i, j = heapq.heappop(heap)
result += current_sum
# 尝试将i+1,j的组合加入堆
if i + 1 < len(array1) and (i + 1, j) not in visited:
heapq.heappush(heap, (array1[i + 1] + array2[j], i + 1, j))
visited.add((i + 1, j))
# 尝试将i,j+1的组合加入堆
if j + 1 < len(array2) and (i, j + 1) not in visited:
heapq.heappush(heap, (array1[i] + array2[j + 1], i, j + 1))
visited.add((i, j + 1))
return result
# 读取输入
array1 = list(map(int, input().split()))
size1 = array1[0]
array1 = array1[1:size1 + 1]
array2 = list(map(int, input().split()))
size2 = array2[0]
array2 = array2[1:size2 + 1]
k = int(input())
# 计算并输出结果
print(min_k_pair_sum(array1, array2, k))
代码解释
- 优先队列初始化:使用最小堆来存储当前可能的最小和组合。初始时,堆中包含array1和array2的第一个元素的和。
- 访问记录:使用集合visited来记录已经处理过的索引对,防止重复计算。
- 处理k次:每次从堆中取出最小的和,并将其相邻的索引对(即i+1,j和i,j+1)加入堆中,前提是这些索引对未被访问过且未越界。
- 输入处理:读取输入的两个数组和k值,确保数组大小和元素正确解析。
- 结果输出:调用min_k_pair_sum函数计算k对元素的最小和,并输出结果。