CODEWITHSUNDEEP
python-3.xlinuxnumpycentos8
Heath
Heath

Reputation: 123

Numpy: Selecting Rows based on Multiple Conditions on Some of its Elements

I believe this is not a duplicate question, although there are questions that are fairly close to this one on the website. I would like to isolate a row from a numpy list given a set of conditions for some of its elements. Here is an example, consider the array Z:

>>> Z = [[1,0,3,4], [1,1,3,6], [1,2,3,9], [1,3,4,0], [2,1,4,5]]
>>> Z = np.array(Z)
>>> Z
array([[1, 0, 3, 4],
       [1, 1, 3, 6],
       [1, 2, 3, 9],
       [1, 3, 4, 0],
       [2, 1, 4, 5]])

and say I would like to isolate the row whose first and second element are both 1. The command that executes that should output the row

np.array([[1, 1, 3, 6]])

However, if I follow this popular question, and make an intuitive extension, such as:

Z[Z[:,0] == 1 & Z[:,1] == 1, :]

I get:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

Is there any quick fix to that? I do not want to iterate over my list. I was wondering if there is a quick "numpy" way for it.

Upvotes: 0

Views: 1174

Answers (2)

johnInHome
johnInHome

Reputation: 449

More simple

print (Z[(Z[:,0]==1)&(Z[:,1]==1)])
or
print (Z[(Z[:,0]==1)&(Z[:,1]==1),:])

You got [[1 1 3 6]]

Upvotes: 0

ansev
ansev

Reputation: 30920

Elegant is np.equal

Z[np.equal(Z[:, [0,1]], 1).all(axis=1)]

Or:

Z[np.equal(Z[:,0], 1) & np.equal(Z[:,1], 1)]

Upvotes: 1

Related Questions