如何高效对长numpy数组按预设间隔切片重塑为单个ndarray

核心优化思路 #

完全舍弃原实现中开销极高的逐次pad、concatenate操作，改为预分配最终数组直接填充有效数据，既省掉了中间临时数组的内存开销，也避免了反复复制全量数组的时间损耗。

优化后代码 #

import numpy as np
temp = np.linspace(0, 99, 100)  # 待处理原始数组
ni, nj = np.array([5, 4, 7, 2]), np.array([1, 5, 7, 13])  # 分块形状参数
idx = ni * nj  # 每个分块的元素个数

# 预计算全局参数
max_ni, max_nj = ni.max(), nj.max()
n_blocks = len(ni)
# 预计算每个分块的起始索引，避免循环中重复累加
start_indices = np.concatenate([[0], np.cumsum(idx[:-1])])
# 直接预分配最终数组，默认填充0等价于原pad的constant模式
x = np.zeros((max_ni, max_nj, n_blocks), dtype=temp.dtype)

# 轻量循环填充数据，无临时变量、无额外复制开销
for i in range(n_blocks):
    start = start_indices[i]
    end = start + idx[i]
    # 切片重塑后直接赋值到对应位置，无中间临时数组
    x[:ni[i], :nj[i], i] = temp[start:end].reshape(ni[i], nj[i])

print(x.shape == (7, 13, 4))

方案优势 #

1. 速度提升显著：原实现逐次concatenate的时间复杂度为O(n²)，n为分块数量，本方案仅需一次内存分配+一次轻量循环，时间复杂度为O(n)，数组越大、分块越多优势越明显。
2. 内存占用极低：无temp1临时变量、无pad产生的中间数组、无concatenate产生的数组副本，内存占用仅略大于最终输出数组的大小，完全适配极大型数组处理场景。
3. 可拓展性强：如果原始数组过大无法载入内存，可直接替换为np.memmap映射的磁盘数组，代码无需修改即可运行，进一步降低内存占用。

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