Как использовать метод Python numpy.where()

В Python мы можем использовать функцию numpy.where() для выбора элементов из массива numpy на основе условия.

Кроме того, мы можем выполнять некоторые операции над этими элементами, если условие выполняется.

Давайте посмотрим, как мы можем использовать эту функцию, используя несколько наглядных примеров!

Синтаксис Python numpy.where()

Эта функция принимает массив, подобный numpy (например, массив NumPy целых чисел/логических значений).

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

Например, condition может принимать значение array([[True, True, True]]), которое представляет собой логический массив, похожий на numpy. (По умолчанию NumPy поддерживает только числовые значения, но мы также можем преобразовать их в bool)

Например, если условие — это array([[True, True, False]]), а наш массив — a=ndarray([[1, 2, 3]]), применив условие к массиву (a[:, condition]), мы получим массив ndarray([[1 2]]).

import numpy as np

a = np.arange(10)
print(a[a <= 2]) # Will only capture elements <= 2 and ignore others

Выход

array([0 1 2])

ПРИМЕЧАНИЕ. То же самое условие может быть представлено как <= 2. Это рекомендуемый формат для массива условий, так как очень утомительно записывать его в виде логического массива.

Но что, если мы хотим сохранить размерность результата и не потерять элементы из исходного массива? Для этого мы можем использовать numpy.where().

numpy.where(condition [, x, y])

У нас есть еще два параметра x и y. Что это?

По сути, это говорит о том, что если condition выполняется для некоторого элемента в нашем массиве, новый массив будет выбирать элементы из x.

В противном случае, если оно ложно, будут взяты элементы из y.

При этом наш окончательный выходной массив будет массивом с элементами из x везде, где condition=True, и элементами из y всякий раз, когда condition=Ложь.

Обратите внимание, что хотя x и y являются необязательными, если вы указываете x, вы ДОЛЖНЫ также указать y. Это связано с тем, что в этом случае форма выходного массива должна быть такой же, как и входной массив.

ПРИМЕЧАНИЕ. Та же логика применима как к одномерным, так и к многомерным массивам. В обоих случаях мы фильтруем по условию. Также помните, что формы x, y и condition передаются вместе.

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

Использование Python numpy.where()

Предположим, мы хотим взять только положительные элементы из массива numpy и установить для всех отрицательных элементов значение 0, давайте напишем код, используя numpy.where().

1. Замените элементы на numpy.where()

Здесь мы будем использовать двумерный случайный массив и выводить только положительные элементы.

import numpy as np

# Random initialization of a (2D array)
a = np.random.randn(2, 3)
print(a)

# b will be all elements of a whenever the condition holds true (i.e only positive elements)
# Otherwise, set it as 0
b = np.where(a > 0, a, 0)

print(b)

Возможный результат

[[-1.06455975  0.94589166 -1.94987123]
 [-1.72083344 -0.69813711  1.05448464]]
[[0.         0.94589166 0.        ]
 [0.         0.         1.05448464]]

Как видите, остались только положительные элементы!

2. Использование numpy.where() только с условием

В приведенном выше коде может возникнуть некоторая путаница, поскольку некоторые из вас могут подумать, что более интуитивным способом было бы просто написать условие следующим образом:

import random
import numpy as np

a = np.random.randn(2, 3)
b = np.where(a > 0)
print(b)

Если вы сейчас попробуете запустить приведенный выше код, с этим изменением вы получите такой вывод:

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

Если вы внимательно посмотрите, b теперь представляет собой набор массивов numpy. И каждый массив — это расположение положительного элемента. Что это значит?

Всякий раз, когда мы предоставляем только условие, эта функция фактически эквивалентна np.asarray.nonzero().

В нашем примере np.asarray(a > 0) вернет логический массив после применения условия, а np.nonzero(arr_like) вернет индексы ненулевые элементы arr_like. (См. эту ссылку)

Итак, теперь мы рассмотрим более простой пример, который показывает нам, насколько гибкими мы можем быть с numpy!

import numpy as np

a = np.arange(10)

b = np.where(a < 5, a, a * 10)

print(a)
print(b)

Выход

[0 1 2 3 4 5 6 7 8 9]
[ 0  1  2  3  4 50 60 70 80 90]

Здесь условие a < 5, которое будет массивом типа numpy [True True True True True False False False False False], x — это массив a, а y — это массив a * 10. Таким образом, мы выбираем только из an, если a < 5, и из a * 10, если a > 5.

Итак, это преобразует все элементы>=5 путем умножения на 10. Вот что мы действительно получаем!

Вещание с помощью numpy.where()

Если мы предоставим все массивы condition, x и y, numpy будет транслировать их вместе.

import numpy as np

a = np.arange(12).reshape(3, 4)

b = np.arange(4).reshape(1, 4)

print(a)
print(b)

# Broadcasts (a < 5, a, and b * 10)
# of shape (3, 4), (3, 4) and (1, 4)
c = np.where(a < 5, a, b * 10)

print(c)

Выход

[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]]
[[0 1 2 3]]
[[ 0  1  2  3]
 [ 4 10 20 30]
 [ 0 10 20 30]]

Опять же, здесь вывод выбирается на основе условия, поэтому все элементы, но здесь b транслируются в форму a. (Одно из его измерений имеет только один элемент, поэтому ошибок при трансляции не будет)

Итак, b теперь станет [[0 1 2 3] [0 1 2 3] [0 1 2 3]], и теперь мы можем выбирать элементы даже из этот широковещательный массив.

Таким образом, форма вывода такая же, как форма a.

Заключение

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

Рекомендации

• Документация SciPy по функции Python numpy.where()