Выберите строки numpy.ndarray, где первый номер строки находится внутри некоторого списка

Question

Я ищу короткий читабельный способ выбрать несколько строк в 2D numpy.ndarray, где первое число каждой строки находится в некотором списке.

Пример:

>>> index
[4, 8]

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

Так что в этом случае мне нужно только

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

потому что первые числа этих строк 4 и 8, которые перечислены в index.

В основном я ищу что-то вроде:

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

что, конечно, не работает.

Answer 1

Если вы хотите использовать понимание списка, вот как (хотя isin, несомненно, лучше для этого случая использования):

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

Answer 2

Вы можете использовать np.isin для проверки, а затем индексировать как обычно:

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