Question

如何在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'）可以为单个值执行此操作价值搜索）。

Answer 1

检查一个阵列与另一个阵列的经典方法是调整形状并使用＆＃39; ==＆＃39;：

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，则其行＆＃39;将丢失，例如[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]慢，但我还没有调查原因。

Answer 2

您可以在已排序的数组上使用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)会比O(n k)快log n < k用于线性搜索{{1}}。

numpy搜索数组，用于多个值，并返回其索引

2 个答案: