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 27 · Answer 1 · edited Jun 14 '14 at 12:29

27

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

edited Jun 14 '14 at 12:29

Veedrac

answered Jun 14 '14 at 12:25

Delta_Fore

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

22

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