Выбор строк из 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, который мне не хватает?
3 ответа
test[numpy.logical_or.reduce([test[:,1] == x for x in wanted])]
Результат должен быть быстрее, чем оригинальная версия, так как NumPy делает внутренние циклы вместо Python.
. Следующее решение должно быть быстрее, чем решение Амнона, так как 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's look up в наборах set, которые для этого намного быстрее, чем списки. Он также быстр, потому что использует быстрые циклы Numpy. Вы также получаете оптимизацию оператора в : как только wanted элемент совпадает, остальные элементы не нужно тестировать (в отличие от "логического или" подхода Amnon, если все элементы в wanted не тестировались, несмотря ни на что).
В качестве альтернативы можно использовать следующий однослойный, который также проходит через массив только один раз:
test[numpy.apply_along_axis(lambda x: x[1] in wanted, 1, test)]
Это гораздо медленнее, однако, так как при каждой итерации элемент во втором столбце извлекается (вместо того, чтобы делать это за один проход, как в первом решении данного ответа).
.Это в два раза быстрее, чем вариант Амнона для len(test)=1000:
wanted = (2,4,6)
wanted2 = numpy.expand_dims(wanted, 1)
print test[numpy.any(test[:, 1] == wanted2, 0), :]