python - 查找一个数组大于第二个数组中的元素的索引

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 解决方案的任何想法？这是非常缓慢的。谢谢

标签： pythonnumpy

Answer 1

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