Python中numpy获取满足条件的索引值(np.where 和 np.argwhere)

1、np.where(condition, x, y)

condition：条件
x：满足条件时函数的输出
y：不满足条件时的输出

>>> import numpy as np
>>> x = np.arange(6).reshape(2, 3)
>>> x
array([[0, 1, 2],
       [3, 4, 5]])
>>> np.where(x > 1, True, False)
array([[False, False,  True],
       [ True,  True,  True]])

如果输入的数组都是1-D数组，np.where()等价于：

[xv if c else yv for c, xv, yv in zip(condition, x, y)]

2、np.where(condition)

没有x和y参数，则以元组形式输出满足条件的列表索引。

>>> np.where(x > 1)
(array([0, 1, 1, 1], dtype=int64), array([2, 0, 1, 2], dtype=int64))

3、np.argwhere(condition)

Find the indices of array elements that are non-zero, grouped by element.
和2类似，但输出的直接是满足要求的元素的坐标索引，不是元组。

>>> np.argwhere(x > 1)
array([[0, 2],
       [1, 0],
       [1, 1],
       [1, 2]], dtype=int64)

参考链接：
1、numpy.where说明文档
2、numpy.argwhere说明文档