选择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

您可以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]])