如何在numpy数组中搜索一小组值(未排序,不应更改)?它应该返回这些值的下标。
例如:a = np.array(['d', 'v', 'h', 'r', 'm', 'a']) # in general it will be large
query = np.array(['a', 'v', 'd'])
# Required:
idnx = someNumpyFunction(a, query)
print(indx) # should be [5, 1, 0]
我是numpy的初学者,我找不到合适的方法来同时为多个值执行此任务(我知道np.where(a=='d')可以为单个值搜索完成)。
检查数组的一个经典方法是调整数组的形状并使用'==':
In [250]: arr==query[:,None]
Out[250]:
array([[False, False, False, False, False, True],
[False, True, False, False, False, False],
[ True, False, False, False, False, False]], dtype=bool)
In [251]: np.where(arr==query[:,None])
Out[251]: (array([0, 1, 2]), array([5, 1, 0]))
如果元素query
没有在a
中找到,则其"行"将缺失,例如[0,2]
而不是[0,1,2]
In [261]: np.where(arr==np.array(['a','x','v'],dtype='S')[:,None])
Out[261]: (array([0, 2]), array([5, 1]))
对于这个小示例,它比等效的列表推导式快得多:
np.hstack([(arr==i).nonzero()[0] for i in query])
它比searchsorted
溶液慢一点。(在这个解中,如果没有找到query
元素,则i
是越界的)。
Stefano建议fromiter
。与list的hstack
相比,它节省了一些时间:
In [313]: timeit np.hstack([(arr==i).nonzero()[0] for i in query])10000 loops, best of 3: 49.5 us per loop
In [314]: timeit np.fromiter(((arr==i).nonzero()[0] for i in query), dtype=int, count=len(query))
10000 loops, best of 3: 35.3 us per loop
但是如果缺少一个元素,或者多次出现,则会引发错误。hstack
可以处理变长条目,fromiter
不能。
np.flatnonzero(arr==i)
比().nonzero()[0]
慢,但我还没有研究为什么。
您可以在排序数组上使用np.searchsorted
,然后将返回的索引恢复到原始数组。你可以使用np.argsort
;如:
>>> indx = a.argsort() # indices that would sort the array
>>> i = np.searchsorted(a[indx], query) # indices in the sorted array
>>> indx[i] # indices with respect to the original array
array([5, 1, 0])
如果a
的大小为n
, query
的大小为k
,则这将是O(n log n + k log n)
,如果log n < k
,则线性搜索将比O(n k)
更快。