在Numpy数组中查找“小于最大值”的最大值的索引的最佳方法

Question

我有一个已排序的numpy数组X，还有两个常量k和delta，它们不在X中。我想找到与k最接近且小于等于k的最大值的索引，并且该值必须在k的delta范围内，即我想要：

max {i | k - delta <= X[i] <= k }    (1)

请注意，这个集合可能为空，在这种情况下，我会返回 None 。目前我觉得我的方法并不是最优的，因为它没有利用X在第一步被排序的事实。

# Get the max from the set of indices in X satisfying (1)
idx = np.where((k-delta <= X) * (X <= k))[0].max()

我不确定在这种情况下Numpy有多聪明，因为它并不知道X已经排序，因此(k-delta <= X) * (X <= k))可能比必要的时间更长。请注意，由于我们自己知道数组已排序，因此可以使用.max()。

有没有更优化的方法？

- rwolst

---

你说这个数组已经排好序了。那么满足你条件的元素不是总是在 k 的前一个索引位置吗？这个事实极大地简化了解决方案。 - DeepSpace

是的，但 k 不在 X 中，所以我必须先比较 k 和 X 中的元素 n。 - rwolst

我明白了。需要注意的是 k 不一定在 X 中。你可能需要在问题中澄清这一点。 - DeepSpace

Answer 1

利用排序顺序的一个高效方法是使用np.searchsorted函数。

def largest_within_delta(X, k, delta):
    right_idx = X.searchsorted(k,'right')-1
    if (k - X[right_idx]) <= delta:
        return right_idx
    else:
        return None

覆盖各种场景的示例运行 -

In [216]: X
Out[216]: array([ 8,  9, 33, 35, 36, 37, 44, 45, 71, 81])

In [217]: largest_within_delta(X, 36, 0) # this k is already in array
Out[217]: 4

In [218]: largest_within_delta(X, 36, 1) # shouldn't choose for next one 37
Out[218]: 4    

In [220]: largest_within_delta(X, 40, 3) # Gets 37's index
Out[220]: 5

In [221]: largest_within_delta(X, 40, 2) # Out of 37's reach

运行时测试

In [212]: # Inputs
     ...: X = np.unique(np.random.randint(0,1000000,(10000)))
     ...: k = 50000
     ...: delta = 100
     ...: 

In [213]: %timeit np.where((k-delta <= X) * (X <= k))[0].max()
10000 loops, best of 3: 44.6 µs per loop

In [214]: %timeit largest_within_delta(X, k, delta)
100000 loops, best of 3: 3.22 µs per loop

- Divakar

---

我喜欢这个，只是一个快速的笔记：为了精确回答问题，函数应该返回 right_idx 而不是 X[right_idx]。 在接受之前，我将等待看看是否有任何速度分析。 - rwolst

@rwolst 我已相应地更新了我的帖子。在那里添加了一些时间。 - Divakar

Answer 2

Numpy.argmax 可以用于利用排序列表。

import numpy as np
np.argmax(X <= k) if k-d < np.argmax(X <= k) < k+d else None

- dms_quant

---

请注意，如果存在并列的情况，np.argmax将返回第一个最大值的索引，但更合理的做法似乎是选择最后一个最大值。此外，我认为np.argmax没有利用排序后的列表。 - Jean Paul