Use numpy.argwhere to get the corresponding values in np.array
I would like to use np.argwhere()
to get values in np.array
.
For example:
z = np.arange(9).reshape(3,3)
[[0 1 2]
[3 4 5]
[6 7 8]]
zi = np.argwhere(z % 3 == 0)
[[0 0]
[1 0]
[2 0]]
I want this array: [0, 3, 6]
and did the following:
t = [z[tuple(i)] for i in zi] # -> [0, 3, 6]
I assume there is an easier way.
source to share
Why not just use masking here:
z[z % 3 == 0]
For your sample matrix, this will generate:
>>> z[z % 3 == 0]
array([0, 3, 6])
If you pass in a matrix with the same dimensions with booleans as indices, you get an array with the elements of that matrix, where is the boolean matrix True
.
This will also work more efficiently since you are filtering at the numpy level (whereas list comprehension works at the Python interpreter level).
source to share
Source argwhere
def argwhere(a):
"""
Find the indices of array elements that are non-zero, grouped by element.
...
"""
return transpose(nonzero(a))
np.where
matches np.nonzero
.
In [902]: z=np.arange(9).reshape(3,3)
In [903]: z%3==0
Out[903]:
array([[ True, False, False],
[ True, False, False],
[ True, False, False]], dtype=bool)
In [904]: np.nonzero(z%3==0)
Out[904]: (array([0, 1, 2], dtype=int32), array([0, 0, 0], dtype=int32))
In [905]: np.transpose(np.nonzero(z%3==0))
Out[905]:
array([[0, 0],
[1, 0],
[2, 0]], dtype=int32)
In [906]: z[[0,1,2], [0,0,0]]
Out[906]: array([0, 3, 6])
z[np.nonzero(z%3==0)]
is equivalent to using I,J
as indexing arrays:
In [907]: I,J =np.nonzero(z%3==0)
In [908]: I
Out[908]: array([0, 1, 2], dtype=int32)
In [909]: J
Out[909]: array([0, 0, 0], dtype=int32)
In [910]: z[I,J]
Out[910]: array([0, 3, 6])
source to share