我如何获取包含一个的 numpy 数组的索引

【问题标题】：How i can get the indexes of numpy array that contain one's我如何获取包含一个的 numpy 数组的索引
【发布时间】：2019-09-16 02:49:02
【问题描述】：

如何以优雅的方式获取 numpy 数组中包含 1 的元素的索引？

我试着做一个循环：

indexes = []
for i in range(len(array)):
    if array[i] == 1:
    indexes += [i]

【问题讨论】：

np.where(array==1)

标签： python numpy numpy-ndarray indices

【解决方案1】：

使用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.flatnonzero

【解决方案2】：

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

【讨论】：