Numpy where() on a 2D matrix

Question

I have a matrix like this

t = np.array([[1,2,3,'foo'],  [2,3,4,'bar'],  [5,6,7,'hello'],  [8,9,1,'bar']]) 

I want to get the indices where the rows contain the string 'bar'

In a 1d array

rows = np.where(t == 'bar') 

should give me the indices [0,3] followed by broadcasting:-

results = t[rows] 

should give me the right rows

But I can't figure out how to get it to work with 2d arrays.

Answer 1

You have to slice the array to the col you want to index:

rows = np.where(t[:,3] == 'bar') result = t1[rows] 

This returns:

 [[2,3,4,'bar'],   [8,9,1,'bar']] 

Answer 2

For the general case, where your search string can be in any column, you can do this:

>>> rows, cols = np.where(t == 'bar') >>> t[rows] array([['2', '3', '4', 'bar'],        ['8', '9', '1', 'bar']],       dtype='|S11')