选择 numpy.ndarray 的行，其中第一行号位于某个列表内

我正在寻找一种简短可读的方法来选择 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]

这当然不起作用。

您可以使用np.isin https://docs.scipy.org/doc/numpy-1.15.1/reference/generated/numpy.isin.html检查，然后照常索引：

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