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