Python Numpy数组：高效查找匹配索引

我有两个列表，一个是巨大的（数百万个元素），另一个是几千个。我想做以下几件事

bigArray=[0,1,0,2,3,2,,.....]

smallArray=[0,1,2,3,4]

for i in len(smallArray):
  pts=np.where(bigArray==smallArray[i])
  #Do stuff with pts...

---

上述方法虽然有效，但速度较慢。有没有什么方法可以在不使用C语言编写东西的情况下更有效地执行此操作？

在您的情况下，预排序大数组可能会使您受益匪浅。下面的示例演示了如何将时间从约45秒减少到2秒（在我的笔记本电脑上）（针对一组特定长度的阵列5e6和1e3）。显然，如果阵列大小相差太大，解决方案就不是最优的。例如，对于默认解决方案，复杂性是O（bigN*smallN），但对于我建议的解决方案，复杂性是O（（bigN+smallN）*log（bigN））

输出：

布鲁特42.5278530121

---

非暴力1.57193303108

到目前为止，我认为没有任何必要使用numpy；您可以使用

defaultdict

，前提是您的内存足够，如果观察次数不超过数百万，则应该使用

big_list = [0,1,0,2,3,2,5,6,7,5,6,4,5,3,4,3,5,6,5]
small_list = [0,1,2,3,4]

from collections import defaultdict

dicto = defaultdict(list) #dictionary stores all the relevant coordinates
                          #so you don't have to search for them later

for ind, ele in enumerate(big_list):
    dicto[ele].append(ind)

结果:

>>> for ele in small_list:
...     print dicto[ele]
... 
[0, 2]
[1]
[3, 5]
[4, 13, 15]
[11, 14]

---

这会给您一些速度。

Numpy提供函数Numpy.searchsorted：

例如：

>>> import numpy as np
>>> sorted = np.argsort(big_list)
>>> r = np.searchsorted(big_list, small_list, side='right',sorter=sorted)
>>> l  = np.searchsorted(big_list, small_list, side='left',sorter=sorted)
>>> for b, e in zip(l, r):
...     inds = sorted[b:e]

---

我真的怀疑你在移植到C时会得到多大的加速，因为比较操作和

where

操作很可能已经在C中实现了。这正是我想要的。谢谢。我不完全确定，但是使用

np.searchsorted

而不是对分循环可能有优化的空间。我没有测试对分版本，但在我的快速实验中，这至少比

defaultdic

查找答案快。在我的设置中，增加了2倍多一点。

>>> import numpy as np
>>> sorted = np.argsort(big_list)
>>> r = np.searchsorted(big_list, small_list, side='right',sorter=sorted)
>>> l  = np.searchsorted(big_list, small_list, side='left',sorter=sorted)
>>> for b, e in zip(l, r):
...     inds = sorted[b:e]