where, nonzero, select — NumPy

Коротко

np.where(условие, x, y) — поэлементный «тернарный оператор»: где условие истинно, берётся x, иначе y. Получается новый массив.
x и y могут быть числами или массивами — они транслируются к форме условия.
np.where(условие) с одним аргументом — то же, что np.nonzero: кортеж массивов индексов, по одному на ось.
np.argwhere возвращает те же индексы парами — удобно перебирать координаты.
Для нескольких условий — np.select: первое сработавшее условие определяет значение.

Синтаксис

np.where(cond, x, y)                  # значения
np.where(cond)                        # индексы (как np.nonzero)
np.argwhere(cond)                     # индексы парами, форма (k, ndim)
np.select(condlist, choicelist, default=0)

Параметры

ПараметрЧто делаетПо умолчанию

condбулев массив или условиеобязательный

x, yзначения для True и False; числа или массивы совместимой формыобязательный

condlistсписок условий; проверяются по порядкуобязательный

choicelistсписок значений той же длины, что condlistобязательный

defaultзначение, если ни одно условие не выполнено0

Примеры

Данные для примеровиспользуются во всех примерах ниже

scores = np.array([72, 45, 90, 58, 81, 39])   # баллы за тест

Одно из двух значений

np.where(scores >= 60, "зачёт", "незачёт")

Вывод

array(['зачёт', 'незачёт', 'зачёт', 'незачёт', 'зачёт', 'незачёт'],
      dtype='<U7')

Заменить часть, остальное оставить

np.where(scores < 50, 50, scores)             # поднять до 50, остальное оставить

Вывод

array([72, 50, 90, 58, 81, 50])

Индексы подходящих элементов

idx = np.where(scores < 60)                   # то же, что np.nonzero
print(idx)
scores[idx]

Вывод

(array([1, 3, 5]),)
array([45, 58, 39])

Индексы в двумерном массиве

grid = np.array([[0, 3, 0],
                 [5, 0, 7]])
print(np.nonzero(grid))                       # (номера строк, номера столбцов)
np.argwhere(grid)                             # пары (строка, столбец)

Вывод

(array([0, 1, 1]), array([1, 0, 2]))
array([[0, 1],
       [1, 0],
       [1, 2]])

Несколько условий: select

grades = np.select(
    [scores >= 85, scores >= 70, scores >= 50],
    ["отлично", "хорошо", "удовл."],
    default="неуд.",
)
grades

Вывод

array(['хорошо', 'неуд.', 'отлично', 'удовл.', 'хорошо', 'неуд.'],
      dtype='<U7')

Подводные камни

Обе ветки вычисляются целиком

np.where получает уже готовые x и y: деление на ноль происходит до выбора и даёт предупреждение (и inf в промежуточном массиве). Безопасный способ — параметр where= у самой операции.

Деление в ветке where

counts = np.array([4, 0, 2])
np.where(counts > 0, 100 / counts, 0)         # 100 / 0 всё равно вычисляется

Вывод

RuntimeWarning: divide by zero encountered in divide
array([25.,  0., 50.])

Деление только там, где можно

counts = np.array([4, 0, 2])
np.divide(100, counts, out=np.zeros(3), where=counts > 0)

Вывод

array([25.,  0., 50.])