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

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

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

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

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

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

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

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

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

import numpy as np

a = np.arange(10)
print(a[a <= 2]) # Будет захвачены только элементы <= 2, остальные будут проигнорированы

Вывод

array([0 1 2])

ЗАМЕТКА: Та же самая условная конструкция может быть представлена также как a <= 2. Это рекомендуемый формат для массива условий, так как очень утомительно записывать его в виде логического массива

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

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

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

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

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

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

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

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

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

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

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

1. Замена элементов с помощью numpy.where()

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

import numpy as np

# Случайная инициализация (2D массив)
a = np.random.randn(2, 3)
print(a)

# b будет содержать все элементы a, когда условие будет истинным (т.е. только положительные элементы)
# В противном случае, установите их как 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. Таким образом, мы выбираем из a только в том случае, если 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)

# Вещание (a < 5, a и b * 10)
# формы (3, 4), (3, 4) и (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.

Вывод

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

Ссылки

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