Use numpy.argwhere to obtain the matching values in an np.array

Question

I'd like to use np.argwhere() to obtain the values in an 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 this:

t = [z[tuple(i)] for i in zi] # -> [0, 3, 6]

I assume there is an easier way.

Answer 1

Why not simply 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 a matrix with the same dimensions with booleans as indices, you get an array with the elements of that matrix where the boolean matrix is True.

This will furthermore work more efficient, since you do the filtering at the numpy level (whereas list comprehension works at the Python interpreter level).

Answer 2

Source for argwhere

def argwhere(a):
    """
    Find the indices of array elements that are non-zero, grouped by element.
    ...
    """
    return transpose(nonzero(a))

np.where is the same as 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])