NumPy中将数组插入另一数组所有位置的无循环/高效实现方法

给定两个NumPy数组，例如a = np.array([1,2,3,4])、b = np.array([5,6])，目标是生成将b插入a所有可插入位置的二维数组，预期输出如下：

[[5 6 1 2 3 4]
 [1 5 6 2 3 4]
 [1 2 5 6 3 4]
 [1 2 3 5 6 4]
 [1 2 3 4 5 6]]

无Python层循环的向量化实现 #

可以完全依靠NumPy底层的向量化操作实现，不需要写显式for循环，核心思路是通过广播构造位置掩码，分区域填充值：

import numpy as np

a = np.array([1,2,3,4])
b = np.array([5,6])
n, m = len(a), len(b)
total_col = n + m
total_row = n + 1

# 构造行列索引矩阵
col_idx = np.arange(total_col)
insert_pos = np.arange(total_row)[:, None]

# 标记三个填充区域
mask_b = (col_idx >= insert_pos) & (col_idx < insert_pos + m)
mask_a_pre = col_idx < insert_pos
mask_a_post = col_idx >= insert_pos + m

# 填充结果
res = np.empty((total_row, total_col), dtype=a.dtype)
res[mask_b] = np.tile(b, total_row)
res[mask_a_pre] = a[col_idx[mask_a_pre]]
res[mask_a_post] = a[col_idx[mask_a_post] - m]

注意：该方案虽然没有Python层循环，但会生成尺寸为(n+1)*(n+m)的布尔掩码数组，当a长度较大时（如超过1万），内存占用会快速升高，不适合超大数组场景。

大长度a场景下的最高效实现 #

当a长度很大时，带递推逻辑的轻量循环反而是最优方案，相比纯向量化实现优势明显：

• 内存开销极低，仅需要存储最终结果数组，无额外的大尺寸掩码、索引数组占用，不会出现内存溢出问题
• 核心操作是连续内存块拷贝，NumPy对这类操作做了极致的SIMD优化，执行速度远高于掩码式向量化实现
• 当b长度固定时，整体时间复杂度和结果数组尺寸线性相关，没有额外计算开销

实现代码如下：

import numpy as np

a = np.array([1,2,3,4])
b = np.array([5,6])
n, m = len(a), len(b)
total_col = n + m
total_row = n + 1

res = np.empty((total_row, total_col), dtype=a.dtype)
# 初始化第一行
res[0, :m] = b
res[0, m:] = a
# 逐行递推生成后续行，仅需少量赋值操作
for i in range(1, total_row):
    res[i] = res[i-1]  # 拷贝上一行的连续内存块，速度极快
    res[i, i-1] = res[i-1, i + m - 1]  # 替换左移出来的a元素位置
    res[i, i:i+m] = res[i-1, i-1:i+m-1]  # 将b块整体右移一位

注：原基础循环示例存在一处笔误：最后一行print(rows.astype(int))应改为print(res.astype(int))才能输出正确的结果数组。

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