Conditional indexing with Numpy ndarray

Question

I have a Numpy ndarray matrix of float values and I need to select spesific rows where certain columns have values satisfying certain criteria. For example lets say I have the following numpy matrix:

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)

Lets say I want to select rows from the matrix where the first column's value is between 1 and 6 and the value of second column is between 2-7.

How can I get the row-indexes of the matrix where these conditions are satisfied? What about if I want to delete the rows that satisfy the conditional criterion?

Solution 1

For a numpy based solution, you can use numpy.where and then get the row indexes from it and then use it for indexing you matrix. Example -

matrix[np.where((1 <= matrix[:,0]) & (matrix[:,0] <= 6)
       & (2 <= matrix[:,1]) & (matrix[:,1] <= 7))]

Demo -

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.]])

---

Another method , as indicated in the comments would be to use boolean masks, Example -

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

matrix[mask,:]

Demo -

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.]])

Solution 2

You can get the indices with :

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

Then newMatrix = np.delete(matrix, rows, axis = 0)

Solution 3

You mentioned MATLAB. Here's the equivalent to the accepted answer using 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

The accepted answer without where is:

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.]])

I switched rows and columns in the Octave version because that is its natural way of generating the same numbers (MATLAB/Octave use the equivalent of numpys 'F' order - see below).

The other changes are 0 v 1 start index, and () v []. Otherwise the two notations are similar.

A simpler way to generate the matrix in 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]])

Or with the MATLAB layout:

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]])

Comments

[jjepsuomi]

Thank you very much! =) Appreciate it

[P. Camilleri]

The part "in range(1, 7)" looks terribly inefficient

[corinna]

There might be even a more elegant method using numpy.where() ...still thinking about it.

[jjepsuomi]

Thank you very much! I'm used to more into Matlab's indexing so Python's indexing sometimes confuses me ;D

[jjepsuomi]

Thank you, excellent! P.S. did you forget to add the 'rows' to the second line?

[Anand S Kumar]

Glad I could be helpful! :-) . If you find the answers helpful, I would like to request you to accept an answer (whichever you find best) by clicking on the tick mark on the left side of the answer, it would be helpful for the community.

[hpaulj]

You could omit where and just use the boolean mask to select the rows. mask=(1<=matrix...); matrix[mask,:].

[Anand S Kumar]

@hpaulj Good suggestion, thank you. I have included that in the answer.

[jjepsuomi]

Thank you very much! =) Excellent!