python - 在 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 . :-)