Python在数组中进行多次搜索

时间:2012-10-19 09:04:25

标签: python arrays numpy

idtopick是一组ids

     idtopick=array([50,48,12,125,3458,155,299,6,7,84,58,63,0,8,-1])

idtolook是包含我感兴趣的ID的另一个数组

     idtolook=array([0,8,12,50])

我想在另一个数组中存储与idtopick对应的idtolook的位置。

这是我的解决方案

    positions=array([where(idtopick==dummy)[0][0] for dummy in idtolook])

导致

    array([12, 13,  2,  0])

它可以工作,但实际上我正在使用的数组存储了数百万点,所以上面的脚本相当慢。我想知道是否有办法让它更快。此外,我希望保持idtolook的顺序,因此任何对其进行排序的算法都不适合我的情况。

1 个答案:

答案 0 :(得分:3)

您可以使用排序:

 sorter = np.argsort(idtopick, kind='mergesort') # you need stable sorting
 sorted_ids = idtopick[sorter]
 positions = np.searchsorted(sorted_ids, idtolook)
 positions = sorter[positions]

请注意,如果idtolook中缺少idtopick,则不会引发错误。你实际上也可以将idtolook排序到结果数组中,这应该更快:

 c = np.concatenate((idtopick, idtolook))
 sorter = np.argsort(c, kind='mergesort')
 #reverse = np.argsort(sorter) # The next two lines are this, but faster:
 reverse = np.empty_like(sorter)
 reverse[sorter] = np.arange(len(sorter))
 positions = sorter[reverse[-len(idtolook):]-1]

与设定的操作有相似之处。