Извлечь значения, которые удовлетворяют условию из массива numpy

Question

Скажите, у меня есть следующие массивы:

a = np.array([1,1,1,2,2,2])
b = np.array([4,6,1,8,2,1])

Можно ли сделать следующее:

a[np.where(b>3)[0]]
#array([1, 1, 2])

Таким образом, выбрать значения из a в соответствии с индексами, в которых выполняется условие в b, но с использованием исключительно np.where или аналогичной функции numpy?

Другими словами, можно ли использовать np.where, указав только массив, из которого можно получить значения при условии True? Или есть другая функция numpy, которая делает это за один шаг?

Answer 1

Давайте воспользуемся пониманием списка, чтобы решить эту проблему -

a = np.array([1,1,1,2,2,2])
b = np.array([4,6,1,8,2,1])
indices = [i for i in range(len(b)) if b[i]>3]  # Returns indexes of b where b > 3 - [0, 1, 3]
a[indices]
    array([1, 1, 2])

Answer 2

Да, есть функция: numpy.extract(condition, array) возвращает все значения из array, которые удовлетворяют условию.

Использование этой функции не слишком эффективно по сравнению с np.where или логическим индексированием. Все эти подходы создают временный логический массив, в котором хранится результат b>3. np.where создает дополнительный массив индексов, тогда как a[b>3] и np.extract напрямую используют логический массив.

Лично я бы использовал a[b>3], потому что это самая краткая форма.

Answer 3

Просто используйте логическое индексирование.

>>> a = np.array([1,1,1,2,2,2])                                                                                                   
>>> b = np.array([4,6,1,8,2,1])                                                                                                   
>>>                                                                                                                               
>>> a[b > 3]                                                                                                                      
array([1, 1, 2])

b > 3 даст вам array([True, True, False, True, False, False]) и с помощью a[b > 3] вы выберете все элементы из a, где индексный массив равен True.