Numpy Where() on a 2D Matrix
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...
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.
2 Answers
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']]
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')