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

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 作为一个集合，但仅构建 bool 数组，而不是一次一行地重建整个数组。对于大型 a 和小型 b，我认为 np.in1d 可能更快。

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

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