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

Question

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.

Answer 1

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