numpy.insert替代方案的向量化实现：按插入后索引插入零

我明白你需要的是和numpy.insert逻辑不同的插入函数——numpy.insert是基于原数组的索引插入元素（所有插入位置都指代插入前的数组），而你想要的是基于插入后的最终数组索引来指定零的位置。下面我来分享一个简单高效的实现方案：

核心思路 #

我们的目标是在最终数组的指定位置放置零，剩下的位置按原数组顺序填充元素。具体步骤如下：

• 先确定最终数组的长度：原数组长度 + 需要插入的零的数量
• 对输入的插入索引进行去重和排序（避免重复索引或乱序导致填充错误）
• 创建一个全零的结果数组，长度为最终数组长度
• 生成掩码标记出不需要放零的位置，将原数组元素填充到这些位置上

代码实现 #

import numpy as np

def insert_zero_post_index(arr, indices):
    # 对索引去重并排序，确保后续逻辑正确
    sorted_indices = np.unique(np.sort(indices))
    original_len = len(arr)
    num_zeros = len(sorted_indices)
    final_length = original_len + num_zeros
    
    # 验证索引是否合法：不能超出最终数组的索引范围
    if np.max(sorted_indices) >= final_length or np.min(sorted_indices) < 0:
        raise ValueError(f"所有索引必须在 0 到 {final_length-1} 之间（最终数组长度为 {final_length}）")
    
    # 创建全零的结果数组，保持和原数组相同的数据类型
    result = np.zeros(final_length, dtype=arr.dtype)
    
    # 生成掩码：标记出不需要放零的位置（True表示填充原数组元素）
    position_mask = np.ones(final_length, dtype=bool)
    position_mask[sorted_indices] = False
    
    # 将原数组元素填充到非零位置
    result[position_mask] = arr
    
    return result

示例演示 #

示例1：在最终数组的索引0、2、4位置插入零 #

arr = np.array([1, 2, 3, 4])
insert_indices = [0, 2, 4]
output = insert_zero_post_index(arr, insert_indices)
print(output)
# 输出：array([0, 1, 0, 2, 0, 3, 4])

示例2：在最终数组的索引1、3、5位置插入零 #

arr = np.array([1, 2, 3, 4])
insert_indices = [1, 3, 5]
output = insert_zero_post_index(arr, insert_indices)
print(output)
# 输出：array([1, 0, 2, 0, 3, 0, 4])

示例3：自动处理重复索引 #

arr = np.array([1, 2, 3, 4])
insert_indices = [2, 2, 4]  # 包含重复索引
output = insert_zero_post_index(arr, insert_indices)
print(output)
# 输出：array([1, 2, 0, 3, 0, 4])

和numpy.insert的区别对比 #

比如用numpy.insert尝试类似操作：

# numpy.insert的索引是基于原数组的，这里插入3个零会报错（原数组只有0-3索引）
numpy.insert(arr, [0,2,4], 0)
# 报错：IndexError: index 4 is out of bounds for axis 0 with size 4

而我们的函数基于最终数组索引，只要索引在最终长度范围内就可以正常工作，完全符合你的需求。

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