查找两个 numpy 数组之间的索引映射

Question

numpy 中是否有一种很好的方法来获取array1 中每个元素在array2 中的位置的元素索引？

一个例子：

array1 = np.array([1, 3, 4])
array2 = np.arange(-2, 5, 1, dtype=np.int)

np.where(array1[0] == array2)
# (array([3]),)
np.where(array1[1] == array2)
# (array([5]),)
np.where(array1[2] == array2)
# (array([6]),)

我想做

np.where(array1 == array2)
# (array([3 5 6]),)

这样的事情可能吗？我们保证array1 中的所有条目都可以在array2 中找到。

Answer 1

方法 #1： 使用 np.in1d 获取匹配位置的掩码，然后使用 np.where 获取这些索引位置 -

np.where(np.in1d(array2, array1))

方法 #2： 使用 np.searchsorted -

np.searchsorted(array2, array1)

请注意，如果array2 没有排序，我们需要使用附加的可选参数sorter。

示例运行 -

In [14]: array1
Out[14]: array([1, 3, 4])

In [15]: array2
Out[15]: array([-2, -1,  0,  1,  2,  3,  4])

In [16]: np.where(np.in1d(array2, array1))
Out[16]: (array([3, 5, 6]),)

In [17]: np.searchsorted(array2, array1)
Out[17]: array([3, 5, 6])

运行时测试-

In [62]: array1 = np.random.choice(10000,1000,replace=0)

In [63]: array2 = np.sort(np.random.choice(100000,10000,replace=0))

In [64]: %timeit np.where(np.in1d(array2, array1))
1000 loops, best of 3: 483 µs per loop

In [65]: %timeit np.searchsorted(array2, array1)
10000 loops, best of 3: 40 µs per loop