numpy.where

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

Возвращает элементы, выбранные из x или y в зависимости от condition.

Примечание

Когда предоставлен только condition, эта функция является сокращенным обозначением для np.asarray(condition).nonzero(). Использование nonzero напрямую предпочтительнее, так как оно корректно работает для подклассов. Остальная часть документации охватывает только случай, когда предоставлены все три аргумента.

Параметры

conditionarray_like, bool

Где True, возвращает x, в противном случае возвращает y.

x, yarray_like

Значения, из которых выбирать. x, y и condition должны быть совместимы для трансляции в некую форму.

Возвращаемое значение

outndarray

Массив с элементами из x, где condition равно True, и элементами из y в противном случае.

См. также

choose

nonzero

Функция, которая вызывается, когда x и y опущены

Примечания

Если все массивы одномерные, where эквивалентно:

[xv if c else yv
 for c, xv, yv in zip(condition, x, y)]

Примеры

>>> a = np.arange(10)
>>> a
array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
>>> np.where(a < 5, a, 10*a)
array([ 0,  1,  2,  3,  4, 50, 60, 70, 80, 90])

Это также может быть использовано для многомерных массивов:

>>> np.where([[True, False], [True, True]],
...          [[1, 2], [3, 4]],
...          [[9, 8], [7, 6]])
array([[1, 8],
       [3, 4]])

Формы x, y и условия транслируются вместе:

>>> x, y = np.ogrid[:3, :4]
>>> np.where(x < y, x, 10 + y)  # both x and 10+y are broadcast
array([[10,  0,  0,  0],
       [10, 11,  1,  1],
       [10, 11, 12,  2]])

>>> a = np.array([[0, 1, 2],
...               [0, 2, 4],
...               [0, 3, 6]])
>>> np.where(a < 4, a, -1)  # -1 is broadcast
array([[ 0,  1,  2],
       [ 0,  2, -1],
       [ 0,  3, -1]])