如何更高效地将可变大小值打包进NumPy数组？

首先明确一点：不存在时间复杂度低于线性的实现方式。因为你必须处理输入中的每一个值（不管是单值还是列表里的元素），总得把每个数据点都放到目标buffer里，这是理论上的下限——任何算法都绕不开遍历所有需要写入的数据。不过我们可以大幅优化线性实现的常数项开销，让实际运行速度比你当前的Python循环快很多，这在大数据量场景下（比如你提到的h5py数据记录）会带来显著提升。

推荐的高效实现 #

核心思路是用NumPy的底层向量化操作替代Python层面的for循环——NumPy的C级操作能避免Python循环的解释器开销，速度会快一个数量级以上。

方法1：扁平化后直接赋值

利用np.concatenate和np.atleast_1d快速把混合结构的values转成一维数组，再直接写入预先分配的buffer：

import numpy as np

buffer = np.zeros(10)
lengths = [1, 1, 1, 2, 1, 3, 1]
values = [1, 2, 3, [4, 5], 6, [7, 8, 9], 10]

# 扁平化所有值：把单值转成长度为1的数组，再拼接成一维数组
flat_values = np.concatenate([np.atleast_1d(v) for v in values])
# 直接赋值到预分配的buffer
buffer[:len(flat_values)] = flat_values

print(buffer)  # [ 1.  2.  3.  4.  5.  6.  7.  8.  9. 10.]

方法2：预计算索引批量赋值（适合需要严格依赖lengths的场景）

如果你的lengths是权威来源（比如担心values里的列表长度和lengths不匹配），可以先计算每个元素的起始/结束索引，再批量赋值：

import numpy as np

buffer = np.zeros(10)
lengths = [1, 1, 1, 2, 1, 3, 1]
values = [1, 2, 3, [4, 5], 6, [7, 8, 9], 10]

# 计算每个元素的起始索引
starts = np.cumsum([0] + lengths[:-1])
ends = starts + lengths

# 遍历批量赋值（这里的循环次数是values的元素数，远小于总数据点数）
for start, end, val in zip(starts, ends, values):
    buffer[start:end] = val

print(buffer)

这个方法的循环次数是values的元素个数（而非总数据点数），如果values里的长列表很多，这个循环的开销也会比你原来的循环小很多。

性能对比 #

假设我们用10000个元素的values（混合单值和长列表）测试：

• 你的原始Python循环：大约需要几毫秒到几十毫秒（取决于数据量）
• 方法1的扁平化赋值：通常快10~100倍，因为所有数据处理都在NumPy的C层完成，避免了Python循环的每次迭代开销

针对h5py场景的额外建议 #

因为你是用虚拟数据集合并IO操作，打包buffer的性能成为瓶颈，建议：

• 尽量让values里的元素本身就是NumPy数组（而非Python列表），这样np.concatenate的效率会更高
• 如果数据量极大，可以考虑分批次扁平化赋值，避免一次性占用过多内存
• 提前验证lengths和values中元素的长度一致性，避免赋值时出现维度错误

内容的提问来源于stack exchange，提问作者bomber8013