NumPy：按np.array过滤行

Question

I'd like to filter a NumPy 2-d array by checking whether another array contains a column value. How can I do that?

import numpy as np

ar = np.array([[1,2],[3,-5],[6,-15],[10,7]])
another_ar = np.array([1,6])
new_ar = ar[ar[:,0] in another_ar]
print new_ar

I hope to get [[1,2],[6,-15]] but above code prints just [1,2] .

Answer 1

You can use np.where ,but note that as ar[:,0] is a list of first elements if ar you need to loop over it and check for membership :

>>> ar[np.where([i in another_ar for i in ar[:,0]])]
array([[  1,   2],
       [  6, -15]])

Answer 2

Instead of using in , you can use np.in1d to check which values in the first column of ar are also in another_ar and then use the boolean index returned to fetch the rows of ar :

>>> ar[np.in1d(ar[:,0], another_ar)]
array([[  1,   2],
       [  6, -15]])

This is likely to be much faster than using any kind of for loop and testing membership with in .