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

后端 未结 2 1588
名媛妹妹
名媛妹妹 2021-02-08 13:24

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          


        
2条回答
  •  盖世英雄少女心
    2021-02-08 13:57

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

提交回复
热议问题