[英]indices of NumPy arrays intersection
我有兩個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關鍵字對sorted進行搜索:
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))方法,但是對於許多用例而言,它仍然應該非常快。
不知道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]
您可以使用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
聲明:本站的技術帖子網頁,遵循CC BY-SA 4.0協議,如果您需要轉載,請注明本站網址或者原文地址。任何問題請咨詢:yoyou2525@163.com.