Python：获取array满足给定条件的元素的索引

使用np.argwhere

import numpy as np
a = np.array([1,2,3,4,5,6,1,2,3])
b = np.argwhere(a == 1)
print(b)

输出：

[[0]
 [6]]

 

使用np.where

import numpy as np
a = np.array([1,2,3,4,5,6,1,2,3])
b = np.where(a == 1)
print(b)
print(b[0])
print(type(b[0]))

输出：

(array([0, 6], dtype=int64),)
[0 6]
<class 'numpy.ndarray'>