从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

numpy.in1d 是你需要的工具：

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

如果你想要处理的数据量很大，这个方法应该是最快的解决方案。而且，我觉得它的代码也最容易理解。

Answer 2

下面这个解决方案在处理更大的 wanted 时应该会比Amnon的方案更快：

# 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 数组搜索一次（而不是像Amnon的答案那样，可能要搜索 len(wanted) 次）。它还利用了Python内置的快速查找功能，使用 集合 来查找元素，这比用列表要快得多。此外，它还使用了Numpy的快速循环，所以速度也很快。还有一个优化是使用了 in 操作符：一旦找到一个匹配的 wanted 元素，剩下的元素就不需要再测试了（而Amnon的方法是“逻辑或”，无论如何都会测试 wanted 中的所有元素）。

另外，你也可以使用下面这个一行代码，它同样只需遍历数组一次：

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

不过，这个方法要慢得多，因为它在每次迭代时都要提取第二列的元素（而不是像第一个解决方案那样一次性完成）。

Answer 3

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

这个结果应该比原来的版本快，因为是NumPy在处理内部循环，而不是Python。