通过一组值过滤numpy数组的最快方法

我是numpy的新手,我也使用pypy 2.2,它有限的numpy支持(请参阅http://buildbot.pypy.org/numpy-status/latest.html)但我想要做的是过滤一个数组通过一组值(即,如果它包含集合中的值,则保持子阵列).我可以使用列表理解,但我宁愿没有中间列表,因为在较长的数组上它并不快,我不禁想到numpy过滤会更快.

>> a = np.array([[   368,    322, 175238,      2],
       [   430,    382, 121486,      2],
       [   451,    412, 153521,      2],
       [   480,    442, 121468,      2],
       [   517,    475, 109543,      2],
       [   543,    503, 121471,      2],
       [   576,    537, 100566,      2],
       [   607,    567, 121473,      2],
       [   640,    597, 153561,      2]])

>> b = {121486, 153521, 121473}

>> np.array([x for x in a if x[2] in b])

>> array([[   430,    382, 121486,      2],
   [   451,    412, 153521,      2],
   [   607,    567, 121473,      2]])

askewchan.. 5

你可以在一行中完成,但你必须使用list(b)它,所以实际上它可能不会更快:

>>> a[np.in1d(a[:,2], list(b))]
array([[   430,    382, 121486,      2],
       [   451,    412, 153521,      2],
       [   607,    567, 121473,      2]])

它的工作原理是因为它np.in1d告诉你第一项中哪一项在第二项:

>>> np.in1d(a[:,2], list(b))
array([False,  True,  True, False, False, False, False,  True, False], dtype=bool)

对于大a而且b,这可能比你的解决方案更快,因为它仍然b用作一个集合,但是只构建布尔数组而不是一次重建一行整个数组.对于大型a和小型b的,我想np.in1d可能会更快.

ainb = np.array([x in b for x in a[:,2]])
a[ainb]

对于小型a和大型b,您自己的解决方案可能是最快的.