Получение первого индекса элемента в массиве NumPy

Как известно, в Python есть возможность получить первый индекс элемента в списке с помощью метода index(). Например, имеется список xs = [1, 2, 3], и мы хотим узнать, под каким индексом в этом списке находится число 2. Для этого мы вызываем метод index() для нашего списка и передаем ему в качестве аргумента число 2: xs.index(2). В результате мы получим число 1, так как именно под этим индексом находится число 2 в нашем списке.

xs = [1, 2, 3]
xs.index(2)  # вернет: 1

Однако, что делать, если у нас есть массив NumPy и мы хотим найти первый индекс определенного элемента в этом массиве? Вот тут уже метод index() не сработает, так как он присущ только для списков Python.

В NumPy для этой задачи можно воспользоваться функцией numpy.where(). Эта функция возвращает индексы элементов, которые удовлетворяют заданному условию.

import numpy as np

arr = np.array([1, 2, 3])
idx = np.where(arr == 2)

print(idx)  # вернет: (array([1]),)

В примере выше мы ищем индекс числа 2 в массиве arr. Функция where() вернет кортеж, содержащий массив с индексами элементов, удовлетворяющих условию (в данном случае arr == 2). В нашем случае в массиве всего один элемент с индексом 1, который удовлетворяет этому условию.

Но обычно нас интересует только первое вхождение элемента, поэтому можно модифицировать код следующим образом:

import numpy as np

arr = np.array([1, 2, 3, 2])
idx = np.where(arr == 2)

print(idx[0][0])  # вернет: 1

Таким образом, мы можем получить первый индекс элемента в массиве NumPy, аналогично тому, как это делается для списков Python с помощью метода index().