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

Question

我有很长的 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]

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

Answer 1

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

您可以将随机变量“每个元素之前的 0 的数量”建模为 Negative Binomial Distribution，以计算“失败”(0) 的数量，然后再准确地获得 1“成功” , 成功概率 1 - p:

# 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])

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

Python: how to vectorized insertions of values into list/np.array (with given probability)?

python

arrays

numpy

insert

vectorization