I have a numpy array:
import numpy as np
stm = np.array([[0,1],[1,5],[4,5],[3,6],[7,9]])
x = 3.5
I want to find list of indices where: stm[i][0]<= x <=stm[i][1]
and the result should be [1,3].
Is there any way to do this in numpy without having to loop through stm?
You could use boolean masking and np.where():
>>> np.where((stm[:,0] <= x) & (stm[:,1] >= x))
(array([1, 3]),)
As an alternative you could have also used np.argwhere (as suggested by @MSeifert), np.nonzero or np.flatnonzero.
They all behave slightly differently, so it's a good idea to know about all of them.
np.argwhere instead (as proposed by @MSeifert) since I think it's the better choice here.argwhere is better since it can only ever provide an index, and not the value itself, like where does. It gives the same result as long as you don't give the (optional) x,y arguments to where though.This can be easily achieved using boolean masks and np.argwhere:
>>> np.argwhere((stm[:,0] <= x) & (stm[:,1] >= x))
array([[1],
[3]], dtype=int64)
or if you want just the first index:
>>> np.argwhere((stm[:,0] <= x) & (stm[:,1] >= x))[:,0]
array([1, 3], dtype=int64)
a bit shorter :
((stm-x).prod(1)<=0).nonzero()
