'How to return a numpy array of the indices of the first element in each row of a numpy array with a given value?
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?
Solution 1:[1]
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])
Solution 2:[2]
This worked for me:
nonzeros = np.nonzero(input)
u, indices = np.unique(nonzeros[0], return_index=True)
expected = nonzeros[1][indices]
Sources
This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.
Source: Stack Overflow
| Solution | Source |
|---|---|
| Solution 1 | hpaulj |
| Solution 2 | MLev |
