使用Numpy ndarray进行条件索引
问题内容:
我有一个浮点值的Numpy ndarray矩阵,我需要选择特殊的行,其中某些列的值满足某些条件。例如,假设我有以下numpy矩阵:
matrix = np.ndarray([4, 5])
matrix[0,:] = range(1,6)
matrix[1,:] = range(6,11)
matrix[2,:] = range(11,16)
matrix[3,:] = range(16,21)
可以说我要从矩阵中选择行,其中第一列的值在1到6之间,第二列的值在2到7之间。
如何获得满足这些条件的矩阵的行索引?如果要删除满足条件条件的行怎么办?
问题答案:
对于基于numpy的解决方案,您可以使用numpy.where
,然后从中获取行索引,然后将其用于为矩阵建立索引。范例-
matrix[np.where((1 <= matrix[:,0]) & (matrix[:,0] <= 6)
& (2 <= matrix[:,1]) & (matrix[:,1] <= 7))]
演示-
In [169]: matrix
Out[169]:
array([[ 1., 2., 3., 4., 5.],
[ 6., 7., 8., 9., 10.],
[ 11., 12., 13., 14., 15.],
[ 16., 17., 18., 19., 20.]])
In [170]: matrix[np.where((1 <= matrix[:,0]) & (matrix[:,0] <= 6)
.....: & (2 <= matrix[:,1]) & (matrix[:,1] <= 7))]
Out[170]:
array([[ 1., 2., 3., 4., 5.],
[ 6., 7., 8., 9., 10.]])
如注释中所示,另一种方法是使用布尔掩码,例如-
mask = ((1 <= matrix[:,0]) & (matrix[:,0] <= 6)
& (2 <= matrix[:,1]) & (matrix[:,1] <= 7))
matrix[mask,:]
演示-
In [41]: matrix
Out[41]:
array([[ 1., 2., 3., 4., 5.],
[ 6., 7., 8., 9., 10.],
[ 11., 12., 13., 14., 15.],
[ 16., 17., 18., 19., 20.]])
In [42]: mask = ((1 <= matrix[:,0]) & (matrix[:,0] <= 6)
....: & (2 <= matrix[:,1]) & (matrix[:,1] <= 7))
In [43]:
In [43]: matrix[mask,:]
Out[43]:
array([[ 1., 2., 3., 4., 5.],
[ 6., 7., 8., 9., 10.]])