Select rows of numpy.ndarray where the first row number is inside some list

user_4130598 0
Viewed 167

I'm looking for a short readable way to select some rows of an 2D numpy.ndarray, where the first number of each row is in some list.

Example: 

>>> index
[4, 8]

>>> data 
array([[ 0,  1,  2,  3],
      [ 4,  5,  6,  7],
      [ 8,  9, 10, 11],
      [12, 13, 14, 15]])

So in this case i only need 

 array([[ 4,  5,  6,  7],
       [8,  9, 10, 11]])

because the first numbers of these rows are 4 and 8 which are listed in index.

Basically im looking for something like:

data[data[:,0] == i if i in index]

which of course is not working.

indexing numpy python
2 Answers

You can use np.isin to check, then index as usual:

idx = [4, 8]

data = np.array([[ 0,  1,  2,  3],
                 [ 4,  5,  6,  7],
                 [ 8,  9, 10, 11],
                 [12, 13, 14, 15]])

>>> data[np.isin(data[:,0], idx)]
array([[ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
user_66711760

If you want to use a list comprehension, here's how (isin is undoubtably better for this usage case, though):

>>> np.array([e for e in data if e[0] in index])
array([[ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
user_80462320
Related