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.

Just to check, is this actually how you created your array? Note that what you've done gives an array of *strings*. If you want a mix of strings and integers, you'll have a record array and it will behave differently. — Andrew Jaffe, Jun 14 '14 at 15:15

score 29 · Answer 1 · edited Jan 30 '23 at 10:25

29

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

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

This returns:

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

edited Jan 30 '23 at 10:25

Tomerikoo

answered Jun 14 '14 at 12:25

Delta_Fore

score 28 · Accepted Answer · answered Jun 14 '14 at 15:11

28

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')

answered Jun 14 '14 at 15:11

Jaime

2 Answers2