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

问题

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.

回答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).

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

来源：https://stackoverflow.com/questions/45255167/use-numpy-argwhere-to-obtain-the-matching-values-in-an-np-array

标签

python

numpy