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