从NumPy ndarray选择行 [英] Selecting rows from a NumPy ndarray

问题描述

我想根据第二个值从 NumPy 数组中仅选择某些行柱子.例如，此测试数组在第二列中具有1到10的整数.

I want to select only certain rows from a NumPy array based on the value in the second column. For example, this test array has integers from 1 to 10 in the second column.

>>> 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的行，那很容易:

If I wanted only rows where the second value is 4, it is easy:

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

但是，当有多个想要的值时，如何获得相同的结果?

But how do I achieve the same result when there is more than one wanted value?

所需列表可以是任意长度.例如，我可能希望第二列为2、4或6的所有行:

The wanted list can be of arbitrary length. For example, I may want all rows where the second column is either 2, 4 or 6:

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

我想出的唯一方法是使用列表理解，然后将其转换回数组，尽管看起来很复杂，但看起来却很复杂:

The only way I have come up with is to use list comprehension and then convert this back into an array and seems too convoluted, although it works:

>>> 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本身有更好的方法吗?

Is there a better way to do this in NumPy itself that I am missing?

推荐答案

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

由于NumPy会执行内部循环而不是Python，因此结果应比原始版本更快.

The result should be faster than the original version since NumPy's doing the inner loops instead of Python.

这篇关于从NumPy ndarray选择行的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！