使用Numpy ndarray进行条件索引

Question

我有一个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之间的行。

如何获得满足这些条件的行索引？如果我想删除符合条件的行，该怎么办？

- jjepsuomi

Answer 1

若要使用基于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.]])

- Anand S Kumar

---

非常感谢！我习惯于Matlab的索引，所以Python的索引有时会让我感到困惑；D - jjepsuomi

你可以省略 where，直接使用布尔掩码来选择行。mask=(1<=matrix...); matrix[mask,:]。 - hpaulj

@hpaulj 好的建议，谢谢。我已经在答案中包含了它。 - Anand S Kumar

Answer 2

你可以使用“：”获取索引。

rows = np.logical_and(0 < matrix[:, 0], < matrix[:, 0] < 6 ) *  np.logical_and(1 < matrix[:, 1], matrix[:, 1] < 7)

然后newMatrix = np.delete(matrix, rows, axis = 0)。

- P. Camilleri

---

谢谢，很好！另外，你有没有忘记在第二行添加“rows”？ - jjepsuomi

Answer 3

你提到了MATLAB。以下是使用Octave等效于被接受答案的方法。

octave:17> ma=reshape(1:20,5,4)
ma =
    1    6   11   16
    2    7   12   17
    3    8   13   18
    4    9   14   19
    5   10   15   20

octave:18> mask=(1<=ma(1,:))&(ma(1,:)<=6)&(2<=ma(2,:))&(ma(2,:)<=7)
mask =
   1   1   0   0

octave:19> ma(:,mask)
ans =
    1    6
    2    7
    3    8
    4    9
    5   10

不使用where的被接受答案是：

In [592]: mask=(1 <= matrix[:,0]) & (matrix[:,0] <= 6) &(2 <= matrix[:,1]) & (matrix[:,1] <= 7)

In [593]: matrix[mask,:]
Out[593]: 
array([[  1.,   2.,   3.,   4.,   5.],
       [  6.,   7.,   8.,   9.,  10.]])

我在Octave版本中交换了行和列，因为这是它生成相同数字的自然方式（MATLAB/Octave使用与numpy的“F”顺序相当-见下文）。

其他更改是0 v 1启动索引和()v[]。否则，两种符号表示法相似。

在numpy中生成矩阵的更简单方法：

In [594]: np.arange(1,21).reshape(4,5)
Out[594]: 
array([[ 1,  2,  3,  4,  5],
       [ 6,  7,  8,  9, 10],
       [11, 12, 13, 14, 15],
       [16, 17, 18, 19, 20]])

或者使用MATLAB布局：

In [595]: np.arange(1,21).reshape(5,4,order='F')
Out[595]: 
array([[ 1,  6, 11, 16],
       [ 2,  7, 12, 17],
       [ 3,  8, 13, 18],
       [ 4,  9, 14, 19],
       [ 5, 10, 15, 20]])

- hpaulj

---

非常感谢！=）太棒了！ - jjepsuomi

Answer 4

获取行索引：

row_indices = [x for x in range(4) if matrix[x][0] in range(1,7) and matrix[x][1] in range(2,8)]

删除行：

indices = [x for x in range(4) if not( matrix[x][0] in range(1,7) and matrix[x][1] in range(2,8))]

new_matrix = matrix[indices]

- corinna

---

非常感谢！=）感激不尽 - jjepsuomi

“in range(1, 7)”这部分看起来非常低效。 - P. Camilleri

可能会有一种更优雅的方法，使用numpy.where()...还在考虑中。 - corinna