numpy.select

numpy.select(condlist, choicelist, default=0)

Функция select() возвращает массив составленный из элементовдругих массивов, которые выбираются из них в зависимости от указанного условия.

Параметры:
condlist - список логических массивов NumPy.
Список логических массивов, выступающих в роли условий выбора элементов из массивов в choicelist. Если в каком-то логическом массиве встречается значение True то выбирается элемент с соответствующим индексом из соответствующего массива в choicelist. Если значение True принимают несколько условий, то выполняется только первое встреченное в списке.
choicelist - список массивов NumPy.
Список той же длины, что и condlist, который содержит массивы NumPy. Элементы данных массивов будут помещены в результирующий массив в зависимости от условий в condlist.
default - скалярное значение (необязательный параметр).
Значение, вставляемое в результат если все условия равны False, по умолчанию равен 0.
Возвращает:
ndarray - массив NumPy
массив с элементами из массивов в choicelist, которые выбраны в соответствии с условиями в axis1 и condlist.
Смотрите так же:
where, take, choose, compress, diag, diagonal


Примеры

Проще всего продемонстрировать работу данной функции для одномерного массива:

>>> import numpy as np
>>> 
>>> a = np.r_[:10]
>>> a
array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
>>> 
>>> x = a < 3
>>> x
array([ True,  True,  True, False, False, False, False, False, False,
       False])
>>> 
>>> y = a > 6
>>> y
array([False, False, False, False, False, False, False,  True,  True,
        True])
>>> 
>>> 
>>> np.select([x, y], [a - 10, a**2])
array([-10,  -9,  -8,   0,   0,   0,   0,  49,  64,  81])

С другой стороны, массивы вовсе не обязаны быть одномерными:

>>> a = np.arange(24).reshape(6, 4)
>>> a
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11],
       [12, 13, 14, 15],
       [16, 17, 18, 19],
       [20, 21, 22, 23]])
>>> 
>>> np.select([(a > 8) & (a < 15), a > 15], [0*a, a + 100], default = -1)
array([[ -1,  -1,  -1,  -1],
       [ -1,  -1,  -1,  -1],
       [ -1,   0,   0,   0],
       [  0,   0,   0,  -1],
       [116, 117, 118, 119],
       [120, 121, 122, 123]])

Предыдущие примеры были выполнены над одним массивом, но все массивы могут быть разными:

>>> a = np.full((2, 5), 3)
>>> a
array([[3, 3, 3, 3, 3],
       [3, 3, 3, 3, 3]])
>>> 
>>> b = np.full((2, 5), 7)
>>> b
array([[7, 7, 7, 7, 7],
       [7, 7, 7, 7, 7]])
>>> 
>>> c = np.arange(10).reshape(2, 5)
>>> c
array([[0, 1, 2, 3, 4],
       [5, 6, 7, 8, 9]])
>>> 
>>> np.select([c%3 == 0, c%3 == 2], [a, b], default = -1)
array([[ 3, -1,  7,  3, -1],
       [ 7,  3, -1,  7,  3]])