查找x落入numpy数组的行的索引

Question

我有一个numpy数组：

import numpy as np
stm = np.array([[0,1],[1,5],[4,5],[3,6],[7,9]]) 
x = 3.5

我想查找索引列表，其中：stm[i][0]<= x <=stm[i][1] 结果应为[1,3]。

有没有办法在numpy中执行此操作而无需遍历stm？

Answer 1

您可以使用布尔屏蔽和np.where()：

>>> np.where((stm[:,0] <= x) & (stm[:,1] >= x))
(array([1, 3]),)

作为替代方案，您也可以使用np.argwhere（由@MSeifert建议），np.nonzero或np.flatnonzero。 它们的行为略有不同，所以了解所有这些都是个好主意。

Answer 2

使用布尔掩码和np.argwhere：

可以轻松实现

>>> np.argwhere((stm[:,0] <= x) & (stm[:,1] >= x))
array([[1],
       [3]], dtype=int64)

或者如果您只想要第一个索引：

>>> np.argwhere((stm[:,0] <= x) & (stm[:,1] >= x))[:,0]
array([1, 3], dtype=int64)

Answer 3

有点短：

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