python - numpy 在数组中首次出现的引用值数组

Question

给定一个阈值 alpha 和一个 numpy 数组 a，找到第一个索引 i 有多种可能性 arr [i] > α;见Numpy first occurrence of value greater than existing value :

numpy.searchsorted(a, alpha)+1
numpy.argmax(a > alpha)

在我的例子中，alpha 可以是标量或任意形状的数组。我想要一个在两种情况下都适用的函数 get_lowest:

alpha = 1.12
arr = numpy.array([0.0, 1.1, 1.2, 3.0])
get_lowest(arr, alpha)  # 2

alpha = numpy.array(1.12, -0.5, 2.7])
arr = numpy.array([0.0, 1.1, 1.2, 3.0])
get_lowest(arr, alpha)  # [2, 0, 3]

有什么提示吗？

Answer 1

您可以使用广播:

In [9]: arr = array([ 0. ,  1.1,  1.2,  3. ])
In [10]: alpha = array([ 1.12, -0.5 ,  2.7 ])
In [11]: np.argmax(arr > np.atleast_2d(alpha).T, axis=1)
Out[11]: array([2, 0, 3])

要折叠多维数组，您可以使用 np.squeeze，但如果您在第一种情况下想要 Python float ，则可能需要做一些特殊的事情:

def get_lowest(arr, alpha):
    b = np.argmax(arr > np.atleast_2d(alpha).T, axis=1)
    b = np.squeeze(b)
    if np.size(b) == 1:
        return float(b)
    return b