I can select rows from a numpy array where the second element is 7 by using myarray[myarray[:,1]==7]. How can I extend this to select rows where the second element is 7 or 9? E.g. something like myarray[myarray[:,1]==7|==9] (obviously that doesn't work).
Add a comment
|
2 Answers
Use a[(a[:,1] == 7) | (a[:,1] == 9)] for example:
In [6]: a = np.array([[4,7,8], [6,9,0], [4,4,4]])
In [7]: a[(a[:,1] == 7) | (a[:,1] == 9)]
Out[7]:
array([[4, 7, 8],
[6, 9, 0]])
Another option is to use numpy.logical_or
In [15]: a[np.logical_or(a[:, 1] == 7, a[:,1] == 9)]
Out[15]:
array([[4, 7, 8],
[6, 9, 0]])
Comments
You can use np.in1d if you want to elegantly include more elements for selection, as shown in the code and sample run next up.
Code -
select_elements_list = [7,9] # Edit this to include more numbers if needed
row_mask = np.in1d(myarray[:,1],select_elements_list) # mask of valid rows
myarray_out = myarray[row_mask,:] # Output with selected rows based on mask
Thus, essentially would be a one-liner like so -
myarray_out = myarray[np.in1d(myarray[:,1],[7,9]),:]
Sample run -
In [15]: myarray
Out[15]:
array([[8, 7, 7, 8, 8],
[8, 9, 8, 9, 9],
[9, 9, 7, 8, 7],
[7, 8, 8, 7, 8]])
In [16]: myarray[:,1]
Out[16]: array([7, 9, 9, 8])
In [17]: row_mask = np.in1d(myarray[:,1],[7, 9])
In [18]: row_mask
Out[18]: array([ True, True, True, False], dtype=bool)
In [19]: myarray[row_mask,:]
Out[19]:
array([[8, 7, 7, 8, 8],
[8, 9, 8, 9, 9],
[9, 9, 7, 8, 7]])
