如何高效在NumPy二维数组的多个位置插入另一个二维数组？

针对大尺寸NumPy数组，循环调用np.insert会因频繁数组复制导致性能低下；而直接批量插入报错，是因为np.insert不支持在多个位置批量插入多行数组。以下是基于NumPy向量操作的高效实现方式：

实现思路 #

通过数组扩展、拼接与重塑完成操作，全程无Python循环，利用NumPy的底层C优化实现最高效率：

1. 将待插入数组to_insert复制N份（N为原数组X的行数），扩展为形状(N, dim, M)的三维数组；
2. 将原数组X扩展为形状(N, 1, M)的三维数组，让每一行单独成为一个子数组；
3. 把上述两个三维数组在第二个维度拼接，得到形状(N, dim+1, M)的数组，此时每个原数组行前已带上dim行待插入数据；
4. 最后将三维数组重塑为二维数组，得到最终结果。

完整代码示例 #

import numpy as np

# 原数组
X = np.array([
    [1, 2, 3, 4, 5],
    [1, 2, 3, 4, 5],
    [1, 2, 3, 4, 5],
])
dim = 2
to_insert = np.zeros((dim, X.shape[1]))

# 高效实现
# 扩展待插入数组为(N, dim, M)
insert_tiled = np.tile(to_insert[None, ...], (X.shape[0], 1, 1))
# 扩展原数组为(N, 1, M)
X_expanded = X[:, None, :]
# 拼接并重塑
combined = np.concatenate([insert_tiled, X_expanded], axis=1)
result = combined.reshape(-1, X.shape[1])

print(result)

输出结果 #

[[0 0 0 0 0]
 [0 0 0 0 0]
 [1 2 3 4 5]
 [0 0 0 0 0]
 [0 0 0 0 0]
 [1 2 3 4 5]
 [0 0 0 0 0]
 [0 0 0 0 0]
 [1 2 3 4 5]]

性能优势 #

该方法的时间复杂度为O(NM(dim+1))，仅涉及几次数组内存分配与底层复制操作，相比循环插入的O(N²*M)复杂度，在大数组场景下性能提升极为明显。

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