CODEWITHSUNDEEP
pythonarraysnumpy
user2295350
user2295350

Reputation: 303

numpy: how to select rows based on a bunch of criteria

How can I fetch the rows for which the second column equals to 4 or 6?

a = np.array(np.mat('1 2; 3 4; 5 6; 7 4'))
b = [4,6]

Apparently, this does not work:

c = a[a[:,1] in b]

Upvotes: 5

Views: 894

Answers (4)

Jaime
Jaime

Reputation: 67427

The numpythonic way of doing this would be to use in1d, something like:

a[np.in1d(a[:, 1], b)]

Upvotes: 10

techiev2
techiev2

Reputation: 298

[x for x in a if x[1] in (4,6,)] should do it for the case you've mentioned.

Edit: As @georgesl suggests, "an ideal way" is to use an iterator right while constructing the array rather than using the numpy.mat method on the string.

Upvotes: 2

lucasg
lucasg

Reputation: 11002

you can use itertools.ifilter :

from itertools import ifilter
c = np.array([ e for e in ifilter(lambda x: x[1] in b, a) ])
>> array([[3,4],
          [5,6],
          [7,4]])

Upvotes: 1

Saullo G. P. Castro
Saullo G. P. Castro

Reputation: 58885

You can do:

check = np.logical_or(a[:,1]==4, a[:,1]==6)
c = a[check,:]

You can also use | for the logical operator or:

check = (a[:,1]==4) | (a[:,1]==6)

Upvotes: 2

Related Questions