NumPy数组交集的索引

时间:2014-03-05 16:48:44

标签: 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]

这样做的简短方法是什么?可以在这里使用列表理解吗?

3 个答案:

答案 0 :(得分:4)

我注意到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))方法,但对于许多用例来说,它应该仍然非常快。

答案 1 :(得分:3)

不确定numpy是否有方法,但这是一种内置方法,它及时采用O(N):

In [9]: lookup = {v:i for i, v in enumerate(arr2)}

In [10]: [lookup[v] for v in arr1]
Out[10]: [0, 1, 0, 2, 2, 1, 0, 3]

答案 2 :(得分:3)

你可以使用NumPy使用广播这样做,但是如果你的阵列很大,你最终可能会为中间结果分配大量的内存

>>> import numpy as np
>>> arr1, arr2 = np.array(['a','b','a','c','c','b','a','d']), np.array(['a','b','c','d'])
>>> arr1 == arr2[:, None]
array([[ True, False,  True, False, False, False,  True, False],
       [False,  True, False, False, False,  True, False, False],
       [False, False, False,  True,  True, False, False, False],
       [False, False, False, False, False, False, False,  True]], dtype=bool)
>>> (arr1 == arr2[:, None]).argmax(axis=0)
array([0, 1, 0, 2, 2, 1, 0, 3])
>>> 

否则请密切注意arraysetops,以防有人向return_index添加intersect1d参数