sandboxj
sandboxj

Reputation: 1254

Numpy: how to select full row based on some columns

a = np.array([[1.,2.,3.],
          [3.,4.,2.],
          [8.,1.,3.]])
b = [8.,1.]
c = a[np.isclose(a[:,0:2],b)]
print(c)

I want to select full rows in a based on only a few columns. My attempt is above.

It works if I include the last column too in that condition, but I don't care about the last column. How do I select rows with 3 columns, based on a condition on 2?

Upvotes: 0

Views: 48

Answers (1)

Divakar
Divakar

Reputation: 221714

Compare with np.isclose using the sliced version of a and then look for all matches along each row, for which we can use np.all or np.logical_and.reduce. Finally, index into input array for the output.

Hence, two solutions -

a[np.isclose(a[:,:2],b).all(axis=1)]
a[np.logical_and.reduce( np.isclose(a[:,:2],b), axis=1)]

Upvotes: 1

Related Questions