NumPy: вычисления и устройство массивов

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

Содержание курса

np.where — поэлементный условный выбор

np.where(condition, x, y) — это векторная замена конструкции if/else, применённой к каждому элементу массива. Функция проходит по всем позициям: где condition истинно, берёт значение из x; где ложно — из y. Результат — новый массив той же формы.

В этом уроке x и y — либо скаляры, либо массивы ровно той же формы, что condition. Это важная оговорка: broadcasting здесь не рассматривается, и если формы не совпадают, NumPy либо сделает что-то неожиданное, либо выбросит ошибку.

Простейший случай — заменить все отрицательные значения нулём:

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]

Здесь condition — это a > 0, булев массив вида [False, False, True, False, True]. Там, где True, берётся соответствующий элемент a; там, где False, подставляется скаляр 0.

Исходный массив a не изменяется — np.where всегда возвращает новый объект.

Условие можно собрать из нескольких проверок с помощью & и | — тех же операторов, что используются для булевых масок при индексации. Синтаксис остаётся прежним: каждое условие в скобках, между ними & или |:

a = np.array([1, 6, 3, 8, 2])
out = np.where((a >= 3) & (a <= 7), a * 10, -1)
print(out)  # [-1 60 30 -1 -1]

Здесь элементы от 3 до 7 умножаются на 10, все остальные заменяются на -1. Значение 8 не попадает в диапазон [3, 7], поэтому получает -1; значение 1 и 2 — тоже.

Обратите внимание: x и y вычисляются полностью до того, как применяется условие. Если написать np.where(a != 0, 1 / a, 0), деление 1 / a вычислится для всех элементов, включая нули — и RuntimeWarning появится, даже если нули никогда не попадут в результат. Это не баг np.where, а следствие того, как Python вычисляет аргументы функции.