numpy.where

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

Функция where() возвращает элементы, которые могут выбираться из двух массивов в зависимости от условия.

Параметры:
condition - массив NumPy или подобный массиву объект.
Булев массив который определяет критерий выбора элементов: True - выбор элемента из x; False - выбор из y.
x - массив NumPy или подобный массиву объект.
Массив из которого выбираются элементы если соответствующие им элементы в condition равны True.
y - массив NumPy или подобный массиву объект.
Массив из которого выбираются элементы если соответствующие им элементы в condition равны False.
Возвращает:
ndarray - массив NumPy
массив с элементами из x и y, выбранными в соответствии с условием в массиве condition.

Замечание

Все три массива должны быть транслируемы к единой форме или иметь одинаковую форму (смотрите примеры).



Примеры

Рассмотрим примеры для двумерных массивов, но не забывайте, что данная функция работает с массивами любой формы

В самом простом случае, все три массива имеют одинаковую форму, при этом массив condition выступает в роли маски, которая накладывается на массивы x и y, и служит условием выбора элементов:

>>> import numpy as np
>>> 
>>> cond = [[True, False],
...         [False, True]]
>>> 
>>> x = np.zeros((2, 2))
>>> x
array([[0., 0.],
       [0., 0.]])
>>> 
>>> y = np.ones((2, 2))
>>> y
array([[1., 1.],
       [1., 1.]])
>>> 
>>> np.where(cond, x, y)
array([[0., 1.],
       [1., 0.]])
>>> 
>>> np.where(cond, y, x)
array([[1., 0.],
       [0., 1.]])

С другой стороны, массивы вовсе недолжны иметь одинаковую форму, но все они должны быть транслируемы друг по другу. Рассмотрим простой пример:

>>> x = np.arange(4).reshape(4, 1)
>>> x
array([[0],
       [1],
       [2],
       [3]])
>>> 
>>> y = np.arange(6).reshape(1, 6)
>>> y
array([[0, 1, 2, 3, 4, 5]])
>>> 
>>> x.shape, y.shape
((4, 1), (1, 6))

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

>>> f = x + y
>>> f
array([[0, 1, 2, 3, 4, 5],
       [1, 2, 3, 4, 5, 6],
       [2, 3, 4, 5, 6, 7],
       [3, 4, 5, 6, 7, 8]])
>>> 
>>> f.shape
(4, 6)

Как видим, массивы x и y транслируются по форме (4, 6), значит теперь осталось добавить условие (массив condition) той же формы или формы которая может быть транслируема к необходимой форме:

>>> cond = f > 4
>>> cond
array([[False, False, False, False, False,  True],
       [False, False, False, False,  True,  True],
       [False, False, False,  True,  True,  True],
       [False, False,  True,  True,  True,  True]])
>>> 
>>> 
>>> np.where(cond, x + 100, y - 100)
array([[-100,  -99,  -98,  -97,  -96,  100],
       [-100,  -99,  -98,  -97,  101,  101],
       [-100,  -99,  -98,  102,  102,  102],
       [-100,  -99,  103,  103,  103,  103]])
>>> 
>>> 
>>> cond_col = [[True], [False], [True], [False]]
>>> 
>>> np.where(cond_col, x - 100, y)
array([[-100, -100, -100, -100, -100, -100],
       [   0,    1,    2,    3,    4,    5],
       [ -98,  -98,  -98,  -98,  -98,  -98],
       [   0,    1,    2,    3,    4,    5]])
>>> 
>>> 
>>> cond_row = [[True, True, False, False, True, True]]
>>> 
>>> np.where(cond_row, x - 100, y)
array([[-100, -100,    2,    3, -100, -100],
       [ -99,  -99,    2,    3,  -99,  -99],
       [ -98,  -98,    2,    3,  -98,  -98],
       [ -97,  -97,    2,    3,  -97,  -97]])

Ну и как вы уже догадались, для работы данной функции достаточно всего одного массива:

>>> arr = np.random.randint(0, 100, (6, 6))
>>> arr
array([[86, 77, 27, 45, 47, 61],
       [29, 70, 38, 36, 65,  5],
       [88, 78, 58, 52, 10, 73],
       [64, 68, 45,  9, 74, 47],
       [96, 27, 24, 66, 74, 36],
       [76, 22, 52, 55,  7, 38]])
>>> 
>>> np.where((40 < arr) & (arr > 60), arr + 1000, arr*0)
array([[1086, 1077,    0,    0,    0, 1061],
       [   0, 1070,    0,    0, 1065,    0],
       [1088, 1078,    0,    0,    0, 1073],
       [1064, 1068,    0,    0, 1074,    0],
       [1096,    0,    0, 1066, 1074,    0],
       [1076,    0,    0,    0,    0,    0]])

В данном примере мы использовали выражение arr*0 что бы для ложных условий вставить значение 0. Но любое скалярное значение может быть транслировано по массиву любой формы, то вместо массивов может быть указано это самое значение и пример выше может быть переписан как:

>>> np.where((40 < arr) & (arr > 60), arr + 1000, 0)
array([[1086, 1077,    0,    0,    0, 1061],
       [   0, 1070,    0,    0, 1065,    0],
       [1088, 1078,    0,    0,    0, 1073],
       [1064, 1068,    0,    0, 1074,    0],
       [1096,    0,    0, 1066, 1074,    0],
       [1076,    0,    0,    0,    0,    0]])

Или даже так:

>>> np.where((40 < arr) & (arr > 60), 1, -1)
array([[ 1,  1, -1, -1, -1,  1],
       [-1,  1, -1, -1,  1, -1],
       [ 1,  1, -1, -1, -1,  1],
       [ 1,  1, -1, -1,  1, -1],
       [ 1, -1, -1,  1,  1, -1],
       [ 1, -1, -1, -1, -1, -1]])