找出一个数组中大于第二个数组元素的索引

Question

我有两个数组

a = np.array([0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16])
b = np.array([0,5,10,15])

我希望得到一个长度为b的输出数组，其中每个元素b[i]是数组a中第一个不小于b[i]的元素的索引：

out = np.array([0, 5, 10, 15]

一个缓慢的解决方案是：

out = []
for x in b: 
    i = np.argmax( a >= x )
    out.append( i )

而这只是一种较小的速度提升：

out = []
i=0
for x in b: 
    i = np.argmax( a[i:] >= x ) + i
    out.append( i )

有没有纯numpy的解决方案？这个速度太慢了。谢谢。

- kevinkayaks

Answer 1

如果 a 已经排序，你可以使用 a.searchsorted(b)。

- Bi Rico

---

谢谢@Bi Rico！非常好用。在3e6个元素上的速度差异惊人。 - kevinkayaks

除了是 "纯 numpy" 之外，searchsorted 还使用了二分查找算法，因此它随着 a 的大小而扩展得更好。它的时间复杂度是 O(log(N)) 而不是 O(N)。 - Bi Rico