How do I vectorize this loop in numpy?

Question

I have this array:

arr = np.array([3, 7, 4])

And these boolean indices:

cond = np.array([False, True, True])

I want to find the index of the maximum value in the array where the boolean condition is true. So I do:

np.ma.array(arr, mask=~cond).argmax()

Which works and returns 1. But if I had an array of boolean indices:

cond = np.array([[False, True, True], [True, False, True]])

Is there a vectorized/numpy way of iterating through the array of boolean indices to return [1, 2]?

Answer 1

For your special use case of argmax, you may use np.where and set the masked values to negative infinity:

>>> inf = np.iinfo('i8').max
>>> np.where(cond, arr, -inf).argmax(axis=1)
array([1, 2])

alternatively, you can manually broadcast using np.tile:

>>> np.ma.array(np.tile(arr, 2).reshape(2, 3), mask=~cond).argmax(axis=1)
array([1, 2])

Answer 2

So you want a vectorized version of:

In [302]: [np.ma.array(arr,mask=~c).argmax() for c in cond]
Out[302]: [1, 2]

What are the realistic dimensions of cond? If the number of rows is small compared to the columns (or length of arr) an iteration like this is probably not expensive.

https://stackoverflow.com/a/31767220/901925 use of tile looks good. Here I change it slightly:

In [308]: np.ma.array(np.tile(arr,(cond.shape[0],1)),mask=~cond).argmax(axis=1)
Out[308]: array([1, 2], dtype=int32)

As expected, the list comprehension times scale with the rows of cond, while the tiling approach is just a bit slower than a single row case. But with times around 92.7 µs this masked array approach is much slower than arr.argmax(). Masking adds a lot of overhead.

The where version is quite a bit faster

np.where(cond, arr, -100).argmax(1)  # 20 µs

A deleted answer suggested

(arr*cond).argmax(1)   # 8 µs

which is even faster. As proposed it didn't work if there are negative arr values. But it can probably be adjusted to handle those.