Logical iteration over numpy array

Question

Suppose that you have two equal-sized lists. First list contains only zeros and ones, and initial value of second list is equal to some fixed number. Other values of second array depends on same-indexed values of first list. Relation between them is that, if value in first list is equal to 0, same-indexed value of second list is equal to preceding one, in all other cases, is equal to some other value. In order to clarify my question, I've written the code below with help of for loop. What is the way to solve this like problem without for loop?

Code
a = np.array([0, 1, 0, 0, 0, 1, 0, 0, 1])
b = np.zeros_like(a)
b[0] = 5
for i in range(1, a.size):
    if a[i] == 0:
        b[i] = b[i-1]
    else:
        b[i] = np.random.randint(5)

Answer 1

Here's a vectorized approach -

offset = int(a[0]!=0)
N = (np.count_nonzero(a!=0)) - offset # no. of rand num to be generated
rand_num = np.append(5,np.random.randint(0,5,N))
out = rand_num[(a!=0).cumsum() - offset]