I have two 1D NumPy arrays x = [x[0], x[1], ..., x[n-1]] and y = [y[0], y[1], ..., y[n-1]]. The array x is known, and I need to determine the values for array y. For every index in np.arange(n), the value of y[index] depends on x[index] and on x[index 1: ]. My code is this:

import numpy as np

n = 5
q = 0.5
x = np.array([1, 2, 0, 1, 0])
y = np.empty(n, dtype=int)

for index in np.arange(n):
  if (x[index] != 0) and (np.any(x[index   1:] == 0)):
    y[index] = np.random.choice([0,1], 1, p=(1-q, q))
  else:
    y[index] = 0

print(y)

The problem with the for loop is that the size of n in my experiment can become very large. Is there any vectorized way to do this?

CodePudding user response：

1. Randomly generate the array y with the full shape.
2. Generate a bool array indicating where to set zeros.
3. Use np.where to set zeros.

Try this,

import numpy as np

n = 5
q = 0.5
x = np.array([1, 2, 0, 1, 0])

y = np.random.choice([0, 1], n, p=(1-q, q))
condition = (x != 0) & (x[::-1].cumprod() == 0)[::-1] # equivalent to the posted one
y = np.where(condition, y, 0)