A simple and clean way: use np.argwhere
to group the indices by element, rather than dimension as in np.nonzero(a)
(i.e., np.argwhere
returns a row for each non-zero element).
>>> a = np.arange(10)
>>> a
array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
>>> np.argwhere(a>4)
array([[5],
[6],
[7],
[8],
[9]])
np.argwhere(a)
is the same as np.transpose(np.nonzero(a))
.
Note: You cannot use a(np.argwhere(a>4))
to get the corresponding values in a
. The recommended way is to use a[(a>4).astype(bool)]
or a[(a>4) != 0]
rather than a[np.nonzero(a>4)]
as they handle 0-d arrays correctly. See the documentation for more details. As can be seen in the following example, a[(a>4).astype(bool)]
and a[(a>4) != 0]
can be simplified to a[a>4]
.
Another example:
>>> a = np.array([5,-15,-8,-5,10])
>>> a
array([ 5, -15, -8, -5, 10])
>>> a > 4
array([ True, False, False, False, True])
>>> a[a > 4]
array([ 5, 10])
>>> a = np.add.outer(a,a)
>>> a
array([[ 10, -10, -3, 0, 15],
[-10, -30, -23, -20, -5],
[ -3, -23, -16, -13, 2],
[ 0, -20, -13, -10, 5],
[ 15, -5, 2, 5, 20]])
>>> a = np.argwhere(a>4)
>>> a
array([[0, 0],
[0, 4],
[3, 4],
[4, 0],
[4, 3],
[4, 4]])
>>> [print(i,j) for i,j in a]
0 0
0 4
3 4
4 0
4 3
4 4