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

Question

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

我试图做一个循环：

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

标签： pythonnumpynumpy-ndarrayindices

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