CODEWITHSUNDEEP
pythonnumpysearch
HuongOrchid
HuongOrchid

Reputation: 332

Find indices of rows that x fall into in numpy array

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?

Upvotes: 0

Views: 217

Answers (3)

Hannes Ovr&#233;n
Hannes Ovr&#233;n

Reputation: 21831

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.

Upvotes: 4

B. M.
B. M.

Reputation: 18628

a bit shorter :

((stm-x).prod(1)<=0).nonzero()

Upvotes: 0

MSeifert
MSeifert

Reputation: 152647

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)

Upvotes: 1

Related Questions