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]])