python - ndarray 大于数组的值的逐行索引

Question

我有一组形状（X，5）：

M = [[1,2,3,4,5],
     [6,7,8,9,1],
     [2,5,7,8,3]
     ...]

和一个形状数组（X，1）：

n = [[3],
     [7],
     [100],
     ...]

现在我需要为每行获取 M >= n 的第一个索引，如果没有这样的索引，则为 nan 。例如：

np.where([1,2,3,4,5] >= 3)[0][0] # Returns 2
np.searchsorted(np.array([1,2,3,4,5]), 3) # Returns 2

这些示例分别应用于每一行（我可以循环 X 次，因为两个数组的长度都是 X）。

我想知道，有没有办法以多维方式一次获得所有索引的输出？

就像是：

np.where(M>=n)

谢谢

编辑： M 中的值未排序，我仍在寻找第一个索引/出现拟合 M >= n （因此可能未搜索排序）

标签： pythonpython-3.xpandasnumpy

Answer 1

您可以首先检查哪些行索引低于或等于n并用于获取每行argmax的第一个。True对于所有列都为 的行False，我们可以使用np.where将它们设置np.nan为例如：

M = np.array([[1,2,3,4,5],
 [6,7,8,9,1],
 [2,5,7,8,3]])

n = np.array([[3],[7],[100]])

le = n<=M

# array([[False, False,  True,  True,  True],
#        [False,  True,  True,  True, False],
#        [False, False, False, False, False]])

lea = le.argmax(1)
has_any = le[np.arange(len(le)), lea]
np.where(has_any, lea, np.nan)
# array([ 2.,  1., nan])