Python：如何将值向量化插入到 list/np.array 中（以给定的概率）？答案

【问题标题】：Python: how to vectorized insertions of values into list/np.array (with given probability)?Python：如何将值向量化插入到 list/np.array 中（以给定的概率）？
【发布时间】：2021-12-23 07:38:11
【问题描述】：

我有很长的 numpy 数组：

v = np.array([10, 15, 15, 15, 10, 30, 30, 10, 10])

我想在每个元素后插入 0 概率

stop_prob = 0.5

所以结果可能如下所示：

[ 0 10  0  0 15  0  0 15 15 10  0  0 30 30 10 10  0  0]

这是我的代码：

v_new = []
for j in range(len(v)+1):
    choice = np.random.choice([1, 0], p=[1-stop_prob, stop_prob])
    while choice == 0:
        v_new.append(0)
        choice = np.random.choice([1, 0], p=[1-stop_prob, stop_prob])
    if j != len(v):
        v_new.append(v[j])

它可以工作，但是对于非常大的列表（具有数百万个值）需要很长时间。如何矢量化这个算法？

这是我的矢量化尝试：

idx = np.random.choice([1, 0], size=len(v), p=[1-stop_prob, stop_prob])
v = np.insert(v, idx, 0)

但结果不正确：

[ 0  0  0  0  0  0  0  0 10  0 15 15 15 10 30 30 10 10]

它将所有零放在列表的开头

【问题讨论】：

标签： python arrays numpy insert vectorization

【解决方案1】：

如果您将0s 以p = stop_prob 的概率添加到v 的每个元素之前，直到您插入该元素，那么这是一个独立伯努利试验序列。

您可以将随机变量“每个元素之前 0 的数量”建模为 Negative Binomial Distribution，以计算“失败”（0）的数量，然后准确地得到 1“成功”，成功概率为 @987654326 @：

# number of zeros we will prepend to each element
# note: use len(v) + 1 if we want trailing zeros, like the original algorithm
num_zeros = np.random.negative_binomial(1, 1 - stop_prob, len(v))

# indices where we will place the elements of v
idx = np.arange(0, len(num_zeros)) # original indices
idx += np.cumsum(num_zeros) # we make space for the zeros

# we build the final array
# note: use (np.max(idx),) if we want trailing zeros
v_new = np.zeros((np.max(idx) + 1,), dtype = v.dtype)
v_new[idx[:len(v)]] = v

一个运行示例：

>>> num_zeros
array([1, 0, 3, 0, 0, 0, 0, 1, 0])
>>> v_new
array([ 0, 10, 15,  0,  0,  0, 15, 15, 10, 30, 30,  0, 10, 10])

【讨论】：