Question

我对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]])

Answer 1

您可以在一行中执行此操作，但必须使用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，您自己的解决方案可能最快。

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

1 个答案: