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