[英]How i can get the indexes of numpy array that contain one's
如何以优雅的方式获取 numpy 数组中包含 1 的元素的索引?
我试着做一个循环:
indexes = []
for i in range(len(array)):
if array[i] == 1:
indexes += [i]
使用np.where
:
a = np.array([0, 0, 1, 1, 0, 1, 1, 1, 0])
np.where(a)
输出:
(array([2, 3, 5, 6, 7], dtype=int64),)
或np.nonzero
:
a.nonzero()
输出:
(array([2, 3, 5, 6, 7], dtype=int64),)
您还可以索引到np.arange
:
np.arange(len(a))[a.astype(bool)]
输出:
array([2, 3, 5, 6, 7])
numpy.argwhere()
可能是一个完美的工作 API。 此外,我们还必须使用arr.squeeze()
删除单例维度。 下面分两种情况:
如果您的输入是0-1
数组,则:
In [101]: a = np.array([0, 0, 1, 1, 0, 1, 1, 1, 0])
In [102]: np.argwhere(a).squeeze()
Out[102]: array([2, 3, 5, 6, 7])
另一方面,如果您有一个通用数组,则:
In [98]: np.random.seed(23)
In [99]: arr = np.random.randint(0, 5, 10)
In [100]: arr
Out[100]: array([3, 0, 1, 0, 4, 3, 2, 1, 3, 3])
In [106]: np.argwhere(arr == 1).squeeze()
Out[106]: array([2, 7])
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.