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

Я хочу выделить только определенные строки из массива 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, который мне не хватает?

Ответы

Ответ 1

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

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

Ответ 2

Следующее решение должно быть быстрее решения Amnon по мере увеличения 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) times). Он также использует встроенный быстрый элемент Python для поиска в наборах, которые для этого намного быстрее, чем списки. Это также быстро, потому что он использует быстрые петли Numpy. Вы также получаете оптимизацию оператора in: как только элемент wanted совпадает, остальные элементы не должны тестироваться (в отличие от "логического" или "подхода Amnon" были все элементы в wanted проверяются независимо от того, что).

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

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

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

Ответ 3

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

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

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

Ответ 4

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

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