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

Answer 2

如果您想使用列表isin推导式，以下是方法（不过，对于这个用例无疑更好）：>>> np.array([e for e in data if e[0] in index])array([[ 4,  5,  6,  7],       [ 8,  9, 10, 11]])