在Numpy数组中查找多个值

Question

a 和 b 是两个整数的 Numpy 数组。它们都是排好序的，并且没有重复的数字。b 是 a 的一个子集。我需要找到 b 中每个元素在 a 中的索引位置。有没有什么高效的 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。:-)