Выбор строк из NumPy ndarray

Question

Я хочу выбрать только определенные строки из массива NumPy на основе значения во втором столбце. Например, этот тестовый массив содержит целые числа от 1 до 10 во втором столбце.

>>> test = numpy.array([numpy.arange(100), numpy.random.randint(1, 11, 100)]).transpose()
>>> test[:10, :]
array([[ 0,  6],
       [ 1,  7],
       [ 2, 10],
       [ 3,  4],
       [ 4,  1],
       [ 5, 10],
       [ 6,  6],
       [ 7,  4],
       [ 8,  6],
       [ 9,  7]])

Если мне нужны только строки, где второе значение равно 4, это просто:

>>> test[test[:, 1] == 4]
array([[ 3,  4],
       [ 7,  4],
       [16,  4],
       ...
       [81,  4],
       [83,  4],
       [88,  4]])

Но как мне достичь того же результата, когда существует более одного требуемого значения?

Список розысков может быть произвольной длины. Например, мне могут потребоваться все строки, где второй столбец равен 2, 4 или 6:

>>> wanted = [2, 4, 6]

Единственный способ, которым я придумал, - это использовать понимание списка, а затем преобразовать его обратно в массив и, кажется, слишком запутанным, хотя это работает:

>>> test[numpy.array([test[x, 1] in wanted for x in range(len(test))])]
array([[ 0,  6],
       [ 3,  4],
       [ 6,  6],
       ...
       [90,  2],
       [91,  6],
       [92,  2]])

Есть ли лучший способ сделать это в самой NumPy, которого мне не хватает?

Answer 1

Следующее решение должно быть быстрее, чем решение Амнона, так как wanted становится больше:

# Much faster look up than with lists, for larger lists:
wanted_set = set(wanted)

@numpy.vectorize
def selected(elmt): return elmt in wanted_set
# Or: selected = numpy.vectorize(wanted_set.__contains__)

print test[selected(test[:, 1])]

Фактически, он имеет преимущество поиска в массиве test только один раз (вместо целых len(wanted) раз, как в ответе Амнона). Он также использует встроенный в Python быстрый поиск элементов в множествах , что намного быстрее, чем списки. Это также быстро, потому что он использует быстрые петли Numpy. Вы также получаете оптимизацию оператора in: как только элемент wanted совпадает, остальные элементы проверять не нужно (в отличие от «логического или» подхода Амнона, все элементы были в wanted) проверяются несмотря ни на что).

В качестве альтернативы, вы можете использовать следующую однострочную строку, которая также проходит через ваш массив только один раз:

test[numpy.apply_along_axis(lambda x: x[1] in wanted, 1, test)]

Это намного медленнее, хотя, так как это извлекает элемент во втором столбце на каждой итерации (вместо того, чтобы делать это за один проход, как в первом решении этого ответа).

Answer 2

test[numpy.logical_or.reduce([test[:,1] == x for x in wanted])]

Результат должен быть быстрее оригинальной версии, поскольку NumPy выполняет внутренние циклы вместо Python.

Answer 3

numpy.in1d ​​ - это то, что вы ищете:

print test[numpy.in1d(test[:,1], wanted)]

Это должно быть самое быстрое решение, если требуется большое; плюс, он самый читаемый, скажем.

Answer 4

Это в два раза быстрее, чем вариант Амнона для len (тест) = 1000:

wanted = (2,4,6)
wanted2 = numpy.expand_dims(wanted, 1)
print test[numpy.any(test[:, 1] == wanted2, 0), :]