Получение только тех значений, которые заполняют условие в массиве numpy

Должен быть (очень) быстрый и эффективный способ получить только элементы из массива numpy или, что еще интереснее, из его фрагмента. Предположим, у меня есть массив numpy:

import numpy as np
a = np.arange(-10,10)

Теперь, если у меня есть список:

s = [9, 12, 13, 14]

Я могу выбрать элементы из:

a[s]  #array([-1,  2,  3,  4])

Как я могу получить массив (numpy) из элементов из [s], которые выполняют условие, то есть являются положительными (или отрицательными)? Это должно привести к тому, что

np.ifcondition(a[s]>0, a[s])  #array([2,  3,  4])

Это выглядит тривиально, но я не смог найти простого и сжатого выражения. Я уверен, что маски делают, но на самом деле это не выглядит прямо. Однако ни один из них:

a[a[s]>0]
a[s[a[s]>0]]

на самом деле являются хорошим выбором.

Спасибо за помощь.

Ответы

Ответ 1

Как насчет:

In [19]: b = a[s]

In [20]: b[b > 0]
Out[20]: array([2, 3, 4])

Ответ 2

Вы обязательно должны принять ответ unutbu, и это то, что я обычно использую для такого рода ситуаций в numpy. Но в интересах иметь несколько способов делать вещи, имея метод, который работает вне numpy, или если промежуточный массив оскорбительно огромен, я добавлю эту альтернативу:

In [3]: [a[S] for S in s if a[S]>0]
Out[3]: [2, 3, 4]

Опять же, метод unutbu значительно быстрее. Но мне нравится этот метод, потому что он может обобщить еще дальше. Если у вас есть дорогая фанковая функция (например, не индексирование) и вы хотите проверить результат этой функции, вы можете сделать это:

In [5]: [f for S in s for f in [FunkyFunction(a[S])] if f>0]
Out[5]: [2, 3, 4]

Странная часть этого заключается в том, что вы создаете список внутри другого списка, но этот внутренний список содержит только один элемент. В основном, что вы делаете, это сохранение значения переменной f, а затем с использованием этого значения дважды: один раз для проверки значения (f>0) и один раз для использования этого значения в списке, если тест проходит.

Ответ 3

Я знаю, что это старый вопрос, и я не уверен, что в то время, когда его спрашивали, было ли доступно следующее, но поскольку я пришел с подобным вопросом, это то, что я нашел, что работало для я:

np.extract(a[s] > 0, a[s])