在Numpy数组中查找多个值

Question

a 和 b 是两个已排序且不含重复元素的 Numpy 整数数组，其中 b 是 a 的一个子集。我需要找到每个 b 元素在 a 中的索引。是否有一种高效的 Numpy 函数可以帮助我避免使用 Python 循环？

（实际上，这些数组是 pandas.DatetimeIndex 和 Numpy datetime64 类型，但我想这并不影响答案。）

- Yariv

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可以用。 :-)

- NPE

---

我甚至不认为他们需要在这里使用二分查找... 在假设两者都已排序的情况下，您可以轻松地说服自己，这可以在O(N)时间内完成。（考虑归并排序的合并阶段）。有趣的是，如果在这些假设下进行Python实现是否能够超越此方法。 - mgilson

@mgilson：您非常正确，OP的问题可以在O(n)时间内解决。我想说的是，searchsorted()解决了一个更一般的问题，因此不能是O(n)。 - NPE

是啊，我也刚意识到这个问题。很遗憾他们没有一个searchdoublesorted函数 :) - mgilson