在Numpy数组中查找多个值

Question

a并且b是两个Numpy整数数组.它们是有序的,没有重复.b是一个子集a.我需要找到a每个元素的索引b.是否有一个有效的Numpy函数可以帮助,所以我可以避免python循环？

(实际上,数组是pandas.DatetimeIndex和Numpy一样datetime64,但我想它并没有改变答案.)

Answer 1

numpy.searchsorted() 可以用来做到这一点:

In [15]: a = np.array([1, 2, 3, 5, 10, 20, 25])

In [16]: b = np.array([1, 5, 20, 25])

In [17]: a.searchsorted(b)
Out[17]: array([0, 3, 5, 6])

据我所知,它不需要b排序,并使用二进制搜索a.这意味着它是O(n logn)而不是O(n).

如果这还不够好,总会有Cython.:-)