使用numpy.argwhere获取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

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])