Python如何按指定概率向list或np.array中向量化插入值？

原代码错误原因 #

你写的向量化代码用错了np.insert的参数：该函数第二个参数为待插入位置的下标数组，你传入的是0/1随机数组，相当于所有0都往数组头部的0、1位置插入，自然会出现所有0堆在开头的问题。

正确实现思路 #

你的原始循环逻辑等价于：在原数组的n+1个间隙（数组开头、每两个元素之间、数组结尾）分别插入随机个数的0，每个间隙插入的0的数量服从参数为1-stop_prob的几何分布（几何分布对应首次成功前的失败次数，这里"成功"是停止插入0，"失败"是插入1个0，完全匹配你的循环逻辑）。

基于这个逻辑可以写出完全向量化的实现，完全避免循环：

1. 生成每个间隙要插入的0的数量：直接调用numpy自带的几何分布随机生成函数
2. 计算每个原数组元素在新数组中的位置
3. 初始化全0新数组，把原元素放到对应位置即可

完整代码示例 #

import numpy as np

# 原数组示例
v = np.array([10, 15, 15, 15, 10, 30, 30, 10, 10])
stop_prob = 0.5

# 生成每个间隙插入的0的个数：几何分布返回首次成功的试验次数，减1得到插入0的数量
insert_zero_counts = np.random.geometric(p=1 - stop_prob, size=len(v) + 1) - 1
# 计算前缀和，得到累计插入的0的数量
prefix_sum = insert_zero_counts.cumsum()
# 计算每个原元素在新数组中的下标
element_positions = prefix_sum[:-1] + np.arange(len(v))
# 初始化全0新数组
new_v = np.zeros(prefix_sum[-1] + len(v), dtype=v.dtype)
# 赋值原元素
new_v[element_positions] = v

性能说明 #

该实现完全基于numpy向量运算，处理长度为百万级别的数组耗时在毫秒级，相比原循环实现性能提升至少2个数量级。

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