How to return a numpy array of the indices of the first element in each row of a numpy array with a-CodePudding

Given a numpy array of shape (2, 4):

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

I want to return an array of shape (N,) where each element of the array is the index of the first True value:

expected = np.array([1, 2])

Is there an easy way to do this using numpy functions and without resorting to standard loops?

CodePudding user response：

np.max with axis finds the max along the dimension; argmax finds the first max index:

In [42]: arr = np.array([[False, True, False, True], [False, False, True, True]])
In [43]: np.argmax(arr, axis=1)
Out[43]: array([1, 2])

CodePudding user response：

This worked for me:

nonzeros = np.nonzero(input)
u, indices = np.unique(nonzeros[0], return_index=True)
expected = nonzeros[1][indices]