numpy.where

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

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

Примечание

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

Параметры:	condition : array_like, bool Где True, возвращает x, в противном случае возвращает y. x, y : array_like Значения для выбора. x, y и condition должны быть совместимы для трансляции в некоторую форму.
Возвращает:	out : ndarray Массив с элементами из 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]])