Векторная арифметика и where/clip

Замена условного Python-цикла векторной операцией

Когда нужно применить одно и то же действие к каждому элементу массива, первый импульс — написать цикл. Это работает, но в NumPy есть прямые замены, которые делают то же самое быстрее и в одну строку.

Самый простой случай — поэлементная арифметика. Цикл вида:

result = np.empty_like(a)
for i in range(len(a)):
    result[i] = a[i] + b[i]

заменяется на:

result = a + b

Результат идентичен, но NumPy выполняет операцию на уровне скомпилированного кода — без Python-overhead на каждой итерации.

Сложнее, когда в цикле есть ветвление:

result = np.empty_like(a)
for i in range(len(a)):
    result[i] = a[i] if a[i] > 0 else 0

Здесь алгоритм перевода в векторный вид состоит из трёх шагов.

1. Определить условие и записать его как булеву маску или выражение: a > 0.
2. Определить, что подставляется при True и при False: при True — a[i] (то есть сам элемент), при False — 0.
3. Вызвать np.where(condition, val_if_true, val_if_false).

Итого:

import numpy as np

a = np.array([-3, 0, 7, -1, 4])
result = np.where(a > 0, a, 0)
print(result)  # [0 0 7 0 4]

Логика та же, что в цикле, только условие, значение при True и значение при False передаются как три аргумента функции.

Если условие составное — его собирают до вызова np.where или прямо внутри него. Например, «заменить на -1 всё, что выходит за диапазон [3, 7]»:

a = np.array([1, 6, 3, 8, 2])
mask = (a >= 3) & (a <= 7)
result = np.where(mask, a, -1)
print(result)  # [-1  6  3 -1 -1]

Предварительно выделять маску в отдельную переменную необязательно — её можно передать прямо в np.where. Отдельная переменная полезна, если условие сложное и его удобнее читать и отлаживать по частям.

Практический критерий: если тело цикла можно описать как «взять одно из двух значений в зависимости от условия», это np.where. Если тело цикла — арифметическая операция без ветвления, это просто поэлементная операция над массивами.