0

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.

ansev
  • 30,322
  • 5
  • 17
  • 31
Heath
  • 123
  • 1
  • 6
  • 1
    Actually, I just figured out a simple solution, one could do: `mask1 = (Z[:,0] == 1);` `mask2 = (Z[:,1] == 1);` `sel = np.logical_and(mask1, mask2);` The numpy array `sel` contains the rows of `Z` where both of the conditions are true. Then, I can do: `index = np.where(sel == 1)[0][0]` and `index` is then the index of interest. Anyways, I would still like to see if there is any more elegant solution to this :D. – Heath Feb 17 '22 at 15:46

2 Answers2

1

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)]
ansev
  • 30,322
  • 5
  • 17
  • 31
0

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

johnInHome
  • 399
  • 3
  • 4