Building one array out of several. Monte-Carlo method

Question

Suppose we have several arrays (in my case, three: a1, a2, a3). We also have one more array of random uniformly distributed numbers, r. They are all of the same length. I need to build an array a combined of those three initial a1, a2, a3 so that if r[i] satisfies a particular condition 1, we take a1[i] as a[i], if it satisfies condition 2, we take a2[i], and a3[i] with condition 3. The three conditions are mutually exclusive.

I've written a for-loop for this:

a = empty(len(r))
for i in range(len(r)):
    if r[i] <= p1:
        a[i] = a1[i]
    if p1 < r[i] <= (p1 + p2):
        a[i] = a2[i]
    if r[i] > (p1 + p2):
        a[i] = a3[i]  

Here conditions 1, 2, 3 are expressed after each "if ...". p1 and p2 are given numbers. This is part of a Monte-Carlo simulation (arrays a1, a2 and a3 are also random numbers with a given distribution). It is too slow, I need to vectorize it somehow, but don't know, how to go about it. Which is the best way? Thanks a lot!

Answer 1

Use numpy.select, which lets you select from multiple conditions:

conds = [r <= p1, r <= (p1 + p2), r > (p1 + p2)]
choices = [a1, a2, a3]
a = np.select(conds, choices)

In your case, the third condition could actually be any condition that evaluates to True (e.g., r == r).