Numpy变长二维数组补零内存分配错误的解决办法

问题背景 #

我有一个大型二维numpy数组（实际是由变长子数组组成的数组结构），其中每个子数组长度不同，示例如下：

[[1,2],[3,4,5,6],[7,8,9]]

需要将所有短于最长子数组的子数组末尾补零，得到如下结果：

[[1,2,0,0],[3,4,5,6],[7,8,9,0]]

我自己编写了补零函数：

def add_zeroes(arr, limit):
    if len(arr) < limit:
        return np.concatenate([arr, np.zeros(limit-len(arr))])
    else:
        return arr

但将其应用到包含60578个子数组的数组时，触发了内存错误：

MemoryError: Unable to allocate 8.59 MiB for an array with shape (1126400,) and data type float64

设备配置为Core i7处理器、16GB内存的Windows 11，需要更高效、更Pythonic的解决方法。

优化方案 #

1. 预分配内存填充（最低内存开销） #

避免循环中反复创建临时数组，先计算最长子数组长度，一次性预分配全零的目标数组，再逐个填充原始数据：

import numpy as np

# 假设原始数据是变长列表（如果是numpy数组，可先转成list处理）
original_data = [[1,2],[3,4,5,6],[7,8,9]]
max_length = max(len(sub) for sub in original_data)

# 预分配与原数据类型一致的全零数组
result_array = np.zeros((len(original_data), max_length), dtype=np.int64)

# 逐个填充子数组内容
for idx, sub_arr in enumerate(original_data):
    result_array[idx, :len(sub_arr)] = sub_arr

这种方式只创建一次目标数组，填充过程仅修改元素，无额外临时内存消耗，是处理大规模数据的最优选择。

2. 用np.pad结合列表推导（代码更简洁） #

如果追求代码简洁性，可以使用numpy的pad函数配合列表推导，但数据量极大时内存效率略低于预分配方式：

import numpy as np

original_data = [[1,2],[3,4,5,6],[7,8,9]]
max_length = max(len(sub) for sub in original_data)

result_array = np.array([
    np.pad(sub, (0, max_length - len(sub)), mode='constant') 
    for sub in original_data
])

问题根源说明 #

原函数的内存问题在于每次调用np.concatenate都会生成新数组，处理6万+子数组时会产生大量临时对象，导致内存碎片化或峰值占用过高。而预分配内存的方式直接锁定所需空间，从根源上避免了临时数组的内存开销。

内容的提问来源于stack exchange，提问作者Ernesto Lopez Fune