Как найти первый индекс в NumPy без where: ускоряем поиск через argmax

Ищете первый индекс в NumPy? Не используйте np.where — это тратит память. Покажем, как найти первое совпадение через np.argmax быстрее и без лишних аллокаций.

Когда набор данных огромный, даже небольшие неэффективности превращаются в реальные узкие места. Частая ситуация — искать первый элемент, удовлетворяющий условию, с помощью функции where() из NumPy, а затем использовать из результата лишь один элемент. Такой подход порождает большие промежуточные массивы, которые на деле не нужны.

Проблема

Типичный паттерн: построить булеву маску, получить индексы через where(), а потом оставить только один. На больших массивах это тратит и время, и память.

import numpy as np

first_axis_ix = np.where(arr > pct / 100)[0]

На первый взгляд может показаться, что это вернёт первый найденный элемент, но это не так. Возвращается первый координатный массив всех совпадений. Иными словами, вы материализуете массив индексов для каждого True в маске и затем почти всё выбрасываете.

Что происходит на самом деле

np.where(condition) возвращает кортеж массивов индексов — по одному на каждую ось. Индексация [0] берёт первый массив из этого кортежа, а не первое совпадение. Для больших входных данных выделение массива всех подходящих позиций означает высокое потребление памяти и лишнюю работу.

Более простой подход с argmax

Если нужен только первый индекс, где условие истинно, стоит запросить именно его. argmax по булевой маске вернёт индекс первого True. Есть тонкость: если значений True нет совсем, argmax вернёт 0, поэтому нужна дополнительная проверка.

import numpy as np

pos = np.argmax(arr > pct / 100)
if arr[pos] > pct / 100:
    first_hit = pos
else:
    first_hit = None

Так вы избегаете построения полных массивов индексов и получаете один целочисленный индекс при наличии совпадения, либо специальное значение в противном случае.

Почему это важно

На огромных массивах создание больших промежуточных структур «на выброс» обходится дорого. Запрос одного индекса через argmax снижает накладные расходы и уменьшает давление на память. Это также лучше выражает намерение: нужен первый матч, а не полный набор совпадающих позиций. При этом важно убедиться, что именно это место — реальная горячая точка вашего конвейера: профилируйте код, чтобы понять, где тратится время. Если критична максимальная производительность end-to-end, для наилучших результатов рекомендуют использовать Numba.

Вывод

Если цель — «найти первый элемент, удовлетворяющий условию», используйте метод, который напрямую возвращает этот индекс. np.where(...) вместе с [0] создаёт лишние массивы и вообще не даёт первое совпадение. Простой argmax по булевой маске с проверкой результата — компактное, бережное к памяти и понятное решение. Проведите профилирование, чтобы оценить эффект, и рассмотрите Numba, если стремитесь к максимальной пропускной способности.