使用布尔向量在Numpy 2D数组中选择行

时间:2013-04-24 20:27:54

标签: python numpy

我有一个矩阵和一个布尔矢量:

>>>from numpy import *
>>>a = arange(20).reshape(4,5)
array([[ 0,  1,  2,  3,  4],
   [ 5,  6,  7,  8,  9],
   [10, 11, 12, 13, 14],
   [15, 16, 17, 18, 19]])

>>>b = asarray( [1, 1, 0, 1] ).reshape(-1,1)
array([[1],
   [1],
   [0],
   [1]])

现在我想选择此矩阵中所有相应的行,其中向量中的相应索引等于零。

>>>a[b==0]
array([10])

我该怎么做才能返回这个特定的行?

[10, 11, 12, 13, 14]

1 个答案:

答案 0 :(得分:3)

b的形状有点奇怪,但如果你能把它制作成更好的索引,那么它就是一个简单的选择:

idx = b.reshape(a.shape[0])
print a[idx==0,:]

>>> [[10 11 12 13 14]]

您可以将其读作“,选择索引为0的所有行,并为选中的每一行取所有列”。您的预期答案应该是列表列表,因为您要求所有符合条件的行。