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