如何在 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
慢一点解决方案。 (在该解决方案中,如果未找到 i
元素,则 query
超出范围)。Stefano 建议
fromiter
.与 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]
慢,但我还没有研究为什么。
关于arrays - numpy 搜索多个值的数组,并返回它们的索引,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/30684563/