python之NumPy 数组交集的索引

我有两个 NumPy 数组。例如:

arr1 = np.array(['a','b','a','c','c','b','a','d']) 
arr2 = np.array(['a','b','c','d']) 

我的任务是创建 arr2 数组的索引列表，其中 arr1 == arr2。

所需列表的长度应等于len(arr1)。例如，在我的例子中，正确答案是 [0,1,0,2,2,1,0,3]。

做到这一点的捷径是什么？是否可以在这里使用列表理解？

请您参考如下方法:

我注意到 arr2 已排序，这是设计使然吗？如果是这样，您可以这样做:

arr1 = np.array(['a','b','a','c','c','b','a','d']) 
arr2 = np.array(['a','b','c','d']) 
 
arr2.searchsorted(arr1) 
# array([0, 1, 0, 2, 2, 1, 0, 3]) 

正如@JAB 提到的，当 arr2 未排序时，您可以使用 sorter 关键字来搜索已排序:

arr2 = np.array(['d', 'c', 'b', 'a']) 
sorter = arr2.argsort() 
sorter[arr2.searchsorted(arr1, sorter=sorter)] 
# array([3, 2, 3, 1, 1, 2, 3, 0]) 

由于 argsort，这是一个复杂度为 O(N*log(N)) 的方法，但对于许多用例来说它应该仍然非常快。