Is there a numpy.where() equivalent for row-wise operations?

Question

I want to find the index of first occurence of some condition row-wise, such that it returns a vector. I would need something like an axis=0 condition in np.where or the pylab find function, but that is not implemented.

To clarify, imagine I have the following matrix:

d=np.array([[0,  1, 0, 1], [0, 1, 1, 1], [1, 0, 0, 0], [0,0,0,1]])

I want the first occurrence of d==1 row wise.

The result should be [1, 1, 0, 3], but I don't see how to do this with np.where or any other function efficiently.

Answer 1

I think what you're looking for here isn't where, which will return you an array of elements from one of two different arrays depending on the condition, but argmax, which returns you the index of the maximum value—or, for a 2D array, the indices of the maximum value of each row or column.

But you don't want the maximum value, you want the values that are 1, right? Well, d==1 is an array of booleans, and True is greater than False, so:

In [43]: np.argmax(d==1, axis=1)
Out[43]: array([1, 1, 0, 3])