Как я могу получить индексы массива NumPy, которые содержат свои

Question

Как я могу элегантно получить индексы элемента, которые содержат 1 в массиве numpy?

Я попытался сделать цикл:

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

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

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