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].
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]])
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.