如何在numpy数组指定索引位置前插入零元素？

嘿，这个需求我熟！用NumPy的向量化操作就能轻松搞定，不用傻乎乎地循环逐个插入（毕竟循环在NumPy里效率太低啦）。直接给你上代码，再一步步解释：

import numpy as np

def insert_zeros(arr, indices):
    # 先对索引排序，兼容用户传入乱序索引的情况
    sorted_indices = np.sort(indices)
    # 计算新数组的长度：原数组长度 + 要插入的零的数量
    new_length = len(arr) + len(sorted_indices)
    # 创建全零的新数组，保持和原数组一致的数据类型
    new_arr = np.zeros(new_length, dtype=arr.dtype)
    # 生成原数组的索引序列
    original_indices = np.arange(len(arr))
    # 计算每个原元素需要偏移的位置：用searchsorted快速统计每个原索引前有多少个要插入的位置
    offsets = np.searchsorted(sorted_indices, original_indices, side='right')
    # 计算原元素在新数组中的目标位置
    target_positions = original_indices + offsets
    # 把原数组的值赋值到新数组的对应位置
    new_arr[target_positions] = arr
    return new_arr

测试示例验证 #

咱们用你给的例子来跑一遍：

示例1 #

arr = np.array([1, 1, 2])
indices = [1]
print(insert_zeros(arr, indices))
# 输出：array([1, 0, 1, 2])

完全符合预期！

示例2 #

假设原数组是np.array([1,2,3,4])，索引列表[0,1,3]：

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

完美实现了在0、1、3号索引前分别插入零的需求。

关键细节说明 #

• 排序索引：不管用户传入的索引是乱序还是有序，先排序能保证searchsorted的计算正确，比如传入[3,0,1]也能得到和[0,1,3]一样的结果。
• 高效偏移计算：np.searchsorted是对数时间复杂度的操作，比循环统计快得多，尤其适合处理大数组。它的作用是快速算出每个原索引前面有多少个要插入零的位置，从而得到偏移量。
• 边界情况处理：
  • 如果传入空索引列表，函数会直接返回原数组的副本（全零数组赋值后和原数组一致）。
  • 如果传入等于原数组长度的索引（比如len(arr)=3时传入[3]），会在数组末尾插入零。
  • 支持重复索引：比如传入[1,1]，会在同一位置前插入两个零，输出array([1,0,0,1,2])。如果不想允许重复索引，可以在函数开头加一句if len(np.unique(sorted_indices)) != len(sorted_indices): raise ValueError("重复索引不被允许")。

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