Selecting rows depending on conditions for multiple columns

Question

I have an array and I want to select rows where I have some condition on different columns of those rows in Python using NumPy. For example consider this array:

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

Now I want all rows where column 1 is 1 and column 4 is 5.
So the output should be [1,2,3,5], only 1 row in this case.

Asked By: Aashish sharma Sharma

||

Source

Answer 1

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

This gives you the desired output:

[[1, 2, 3, 5]]

For an array like this

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

you would then obtain

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

Assuming you call the result res, you can the easily access the results. e.g. res[1] would then be [1, 98, 76, 5].

If – for whatever reason – you would like to get the two numbers between the 1 and the 5 you could use

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

which would give you

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

Is that what you were looking for?

Answered By: Cleb

Answer 2

You can use following approach :

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

Answered By: Mazdak

Answer 3

Boolean indexing/masking does the job nicely

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

When calculating I, use () freely so the == tests have priority over the &. Kasras test would also work: I = np.all(test_array[:,[0,3]]==[1,5],axis=1).

Answered By: hpaulj

Answer 4

If you want a function to be the criteria for rows, you can do:

def filter(row):
    if row[0]==1 and row[3]==5:
        return True
    return False

out_array = test_array[numpy.array([filter(row) for row in test_array])]

numpy.array([filter(row) for row in test_array]) gives you a boolean array corresponding to which rows you want, and which you don’t.

Answered By: Pallav Agarwal

Answer 5

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

Answered By: Rakesh

Selecting rows depending on conditions for multiple columns

Question:

Answers: