按条件查找数组元素并返回索引——np.where() 和 np.argwhere()

np.where()

返回满足括号内条件的元素的索引，元组形式，注意是（第一维索引，第二维索引，。。。）这样的形式，这样非常方便再次引用。
格式 : np.where(condition, [x, y])
[x,y]为可选项，满足条件返回x，不满足条件返回y。

Example：

一般替换

>>> a = np.array([[0, 1, 2],
...               [0, 2, 4],
...               [0, 3, 6]])
>>> np.where(a < 4, a, -1)  # -1 is broadcast
array([[ 0,  1,  2],
       [ 0,  2, -1],
       [ 0,  3, -1]])

高端一点，注意元素对应的位置。

>>> np.where([[True, False], [True, True]],
...          [[1, 2], [3, 4]],
...          [[9, 8], [7, 6]])
array([[1, 8],
       [3, 4]])

重新引用元素

>>>a = np.random.randint(10, size=(5,4))
>>>a
Out[58]: 
array([[0, 4, 7, 8],
       [8, 6, 4, 9],
       [5, 5, 2, 6],
       [4, 0, 7, 4],
       [9, 4, 3, 6]])
       
>>>np.where(a>5)
Out[59]: 
(array([0, 0, 1, 1, 1, 2, 3, 4, 4], dtype=int64),
 array([2, 3, 0, 1, 3, 3, 2, 0, 3], dtype=int64))
 
重新引用符合要求的元素
>>>a[np.where(a>5)[0], np.where(a>5)[1]]
Out[66]: array([7, 8, 8, 6, 9, 6, 7, 9, 6])

np.argwhere()

返回满足括号内条件的元素的索引，数组形式，但是看起来舒服多了。

>>>a = np.random.randint(10, size=(5,4))
>>>a
Out[58]: 
array([[0, 4, 7, 8],
       [8, 6, 4, 9],
       [5, 5, 2, 6],
       [4, 0, 7, 4],
       [9, 4, 3, 6]])
>>>np.argwhere(a>5)
Out[60]: 
array([[0, 2],
       [0, 3],
       [1, 0],
       [1, 1],
       [1, 3],
       [2, 3],
       [3, 2],
       [4, 0],
       [4, 3]], dtype=int64)