Самый быстрый способ фильтрации массива 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 в качестве набора, но строит только логический массив вместо того, чтобы перестраивать весь массив по одной строке за раз. Для больших a и небольших b, я думаю, np.in1d может быть быстрее.

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

Для малых и больших ab, ваше собственное решение, вероятно, самый быстрый.