Блог

У меня очень длинный массив 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]
 

Он помещает все нули в начало списка

Если вы добавляете 0s к каждому элементу v с вероятностью p = stop_prob , пока не вставите элемент, то это последовательность независимых испытаний Бернулли.

Вы можете смоделировать случайную величину «число 0 перед каждым элементом» как отрицательное биномиальное распределение, чтобы подсчитать количество «сбоев» (0s), прежде чем получить точно 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])