P1631 序列合并 数据较弱
题意： 给两个长度为 \(N\) 的单调不降序列 \(A,B\) ，在 \(A,B\) 中各取一个数可以得到 \(N^2\) 个和，求这 \(N^2\) 个和的最小的 \(N\) 个。
序列一： \(A_1,A_2,A_3\cdots A_N\)
序列二： \(B_1,B_2,B_3\cdots B_N\)
大根堆维护
大根堆维护最小的 \(N\) 个，时间复杂度为 \(O(N^2)\) 数据较弱剪枝可以通过。
- 将 \(A_1\) 与序列二组合初始化大根堆。
- 枚举剩余的 \(A\) 元素和 \(B\) 元素的和不断地与堆顶比较，小于就替换堆顶元素，大于
break。 - 最后堆中的元素就是最小的 \(N\) 个。
思路简单，但实现起来有点麻烦，基础太差了说是（
- Python 的
heapq模块是小跟堆，对已知的数组需要heapify进行堆化，如果从空数组开始heappush的话就不用堆化。 - 不需要弹出时直接使用索引获取堆顶元素 \(hq[0]\) 来进行比较等操作。
heapreplace函数可以很方便的替换堆顶元素。
具体代码为
import sys
from heapq import heappop, heappush, heapify, heapreplace
if 1:inp = lambda: sys.stdin.readline().strip()II = lambda: int(inp())MII = lambda: map(int, inp().split())LII = lambda: list(MII())Max = lambda x, y: x if x > y else yMin = lambda x, y: x if x < y else ydef main():n = II()a = LII()b = LII()hq = [-(a[0] + b[i]) for i in range(n)]heapify(hq) #! 细节三 已知的数组必须要堆化for i in range(1, n):#! 细节二if a[i] + b[0] >= -hq[0]:breakfor j in range(n):val = a[i] + b[j]if val >= -hq[0]:breakelse:heapreplace(hq, -val) #! 细节一print(*sorted([-x for x in hq]))if __name__ == "__main__":main()
小根堆多路合并
让 \(A\) 中每一个数都和 \(B\) 相加得到 \(N\) 个递增的序列，先将第一列放入小根堆中。
- \(A_1+B_1 \le A_1+B_2\le \cdots \le A_1 + B_N\)
- \(A_2+B_1 \le A_2+B_2 \le \cdots \le A_2 + B_N\)
- \(\cdots\)
- \(A_N+B_1 \le A_N+B_2 \le \cdots \le A_N + B_N\)
- 取出小根堆的堆顶元素放入数组，然后将其所在队列的后一位元素放入堆中。因此需要记录下标。
- 这样循环操作 \(N\) 次就可以得到最小的 \(N\) 个元素了。
时间复杂度为 \(N\log N\) 。
具体代码为
# 模版代码直接复制即可
def main():n = II()a = LII()b = LII()hq = []for i, x in enumerate(a):heappush(hq, (b[0] + x, i, 0))ans = []for _ in range(n):cur, i, j = heappop(hq)ans.append(cur)if j + 1 < n: # 防止越界heappush(hq, (a[i] + b[j + 1], i, j + 1))print(*ans)