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

问题描述

我正在寻找一种简短的可读方式来选择2D numpy.ndarray的某些行，其中每行的第一个数字在某个列表中.

I'm looking for a short readable way to select some rows of an 2D numpy.ndarray, where the first number of each row is in some list.

示例:

>>> index
[4, 8]

>>> data
array([[ 0,  1,  2,  3],
      [ 4,  5,  6,  7],
      [ 8,  9, 10, 11],
      [12, 13, 14, 15]])

所以在这种情况下，我只需要

So in this case i only need

 array([[ 4,  5,  6,  7],
       [8,  9, 10, 11]])

因为这些行的第一个数字是index中列出的4和8.

because the first numbers of these rows are 4 and 8 which are listed in index.

基本上，我正在寻找类似的东西:

Basically im looking for something like:

data[data[:,0] == i if i in index]

哪个当然不起作用.

推荐答案

您可以使用 np.isin 进行检查，然后照常进行索引:

You can use np.isin to check, then index as usual:

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

这篇关于选择numpy.ndarray的行，其中第一行号在某个列表内的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持！