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

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

2 个答案: