序列合并

GitHub跳转原题关系图返回列表

每个 A[i] 与整个 B 形成一条有序和序列，用堆归并取最小 n 个。

OJ: luogu

题目 ID: P1631

难度:普及+/提高

标签:多路归并二叉堆python

日期: 2026-07-16 21:00

目录

  1. 题意
  2. 思路
  3. Python 知识
  4. 代码
  5. 复杂度
  6. 总结

题意

两个不降数组产生 n2n^2 个两两和，输出最小的 n 个。

思路

固定 A[i] 后，A[i]+B[0..] 有序。把每行第一项放入堆；弹出某行第 j 项后，只补入同一行第 j+1 项，弹 n 次即可。

Python 知识

  • 堆项 (sum,row,column) 完整记录下一候选来源。
  • enumerate(first) 同时生成行号和行基值。
  • 输入有序，因此无需额外排序。

代码

python
import heapq
import sys


data = iter(map(int, sys.stdin.buffer.read().split()))
n = next(data)
first = [next(data) for _ in range(n)]
second = [next(data) for _ in range(n)]
heap = [(value + second[0], i, 0) for i, value in enumerate(first)]
heapq.heapify(heap)
answers = []

for _ in range(n):
    result, row, column = heapq.heappop(heap)
    answers.append(result)
    if column + 1 < n:
        column += 1
        heapq.heappush(heap, (first[row] + second[column], row, column))
print(*answers)

复杂度

时间 O(nlog⁡n)O(n\log n)，空间 O(n)O(n)。

总结

不要枚举 n2n^2 个和；把矩阵每一行当作一条有序流即可。