Numpy: search for the first matching row

Question

If I have a numpy array like this:

import numpy as np
x = np.array([[0,1],[0,2],[1,1],[0,2]])

How can I return the index of the first row that matches [0,2]?

For lists it's easy using index:

[[0,1],[0,2],[1,1],[0,2]]

l.index([0,2])
> 1

And I know numpy has the function numpy.where, but I'm not really sure what to make of the output of numpy.where:

np.where(x==[0,2])
> (array([0, 1, 1, 3, 3]), array([0, 0, 1, 0, 1]))

There's also numpy.argmax, but that also doesn't return what I'm looking for, which is simply the index 1

np.argmax(x == [0,2], axis = 1)

Divakar · Accepted Answer · 2015-10-26 17:07:37Z

4

If the search list is [0,2], you would bring in broadcasting when compared against x, giving us a mask of the same shape as x. Since you are looking for an exact match, you would look for rows with all TRUE values with .all(1). Finally, you want the first index, so use np.where or np.nonzero and select the first element. The implementations as sample runs would be -

In [132]: x
Out[132]: 
array([[0, 1],
       [0, 2],
       [1, 1],
       [0, 2]])

In [133]: search_list = [0,2]

In [134]: np.where((x == search_list).all(1))[0][0]
Out[134]: 1

In [135]: np.nonzero((x == search_list).all(1))[0][0]
Out[135]: 1

edited Oct 26, 2015 at 17:07

answered Oct 26, 2015 at 16:48

Divakar

222k19 gold badges273 silver badges374 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

Numpy: search for the first matching row

1 Answer 1

Comments

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Related