Выбор строк из 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

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

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

Answer 2

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

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

@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 соответствует, остальные элементы не должны быть протестированы (в отличие от «логического» или «подхода Amnon», были проверены все элементы в wanted независимо от того, что) ,

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

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

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

Answer 3

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

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

Answer 4

numpy.in1d является то, что вы ищете:

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

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