根据多列的条件选择行

Question

我有一个数组，我想使用NumPy选择在这些行的不同列上有一些条件的行。例如，考虑以下数组：

test_array = numpy.array(([1,2,3,5],[4,5,6,7],[7,8,9,4]))

现在，我需要所有行，其中第1列是1，第4列是5。因此，输出应该是[1,2,3,5]，在本例中只有1行。

Answer 1

您可以使用以下方法：

>>> test_array[np.where(np.all(test_array[:,[0,3]]==[1,5],axis=1))]
array([[1, 2, 3, 5]])

Answer 2

[list(x) for x in test_array if x[0]==1 and x[3]==5]

这将为您提供所需的输出：

[[1, 2, 3, 5]]

像这样的数组

test_array=numpy.array(([1,2,3,5],[4,5,6,7],[7,8,9,4],[1,98,76,5]))

然后你会得到

[[1, 2, 3, 5], [1, 98, 76, 5]]

假设您调用结果res，您可以轻松地访问结果。这样，res[1]就是[1, 98, 76, 5]。

如果--不管出于什么原因--你想得到1和5之间的两个数字。

[sl[1:3] for sl in res]

这会让你

[[2, 3], [98, 76]]

这就是你要找的吗？

Answer 3

布尔索引/掩蔽很好地完成了这项工作。

In [498]: test_array=np.array(([1,2,3,5],[4,5,6,7],[7,8,9,4]))

In [499]: I = (test_array[:,0]==1) & (test_array[:,3]==5)

In [500]: I
Out[500]: array([ True, False, False], dtype=bool)

In [501]: test_array[I,:]
Out[501]: array([[1, 2, 3, 5]])

在计算I时，可以自由地使用()，这样==测试就比&具有优先权。Kasra的测试也会起作用：I = np.all(test_array[:,[0,3]]==[1,5],axis=1)。