用想要的行列表切片 numpy 数组

Posted 2023-02-23

【中文标题】用想要的行列表切片 numpy 数组

【英文标题】：Slice numpy array wth list of wanted rows

【发布时间】：2012-11-19 12:04:24

【问题描述】：

我有一个 numpy 二维数组 A 和一个行号列表 row_set。如何获取新数组B，例如如果row_set = [0, 2, 5]，则B = [A_row[0], A_row[2], A_row[5]]？

我想到了这样的事情：

def slice_matrix(A, row_set):
    slice = array([row for row in A if row_num in row_set])

但我不知道如何获得 row_num。

【参考方案1】：

使用take():

In [87]: m = np.random.random((6, 2))

In [88]: m
Out[88]: 
array([[ 0.6641412 ,  0.31556053],
       [ 0.11480163,  0.00143887],
       [ 0.4677745 ,  0.43055324],
       [ 0.49749099,  0.15678506],
       [ 0.48024596,  0.65701218],
       [ 0.48952677,  0.97089177]])

In [89]: m.take([0, 2, 5], axis=0)
Out[89]: 
array([[ 0.6641412 ,  0.31556053],
       [ 0.4677745 ,  0.43055324],
       [ 0.48952677,  0.97089177]])

【参考方案2】：

您可以将列表或数组作为索引传递给任何 np 数组。

>>> r = np.random.randint(0,10,(5,5))
>>> r
array([[3, 8, 9, 8, 4],
       [4, 1, 5, 9, 1],
       [3, 6, 8, 8, 0],
       [5, 1, 7, 6, 1],
       [6, 1, 7, 7, 7]])
>>> idx = [0,3,1]
>>> r[idx]
array([[3, 8, 9, 8, 4],
       [5, 1, 7, 6, 1],
       [4, 1, 5, 9, 1]])

【参考方案3】：

速度对比：take() 更快。

In [1]:  m = np.random.random((1000, 2))
         i = np.random.randint(1000, size=500)

         %timeit m[i]
Out[1]:
         10000 loops, best of 3: 27.2 µs per loop

In [2]:  %timeit m.take(i, axis=0)
Out[2]:
         100000 loops, best of 3: 7.24 µs per loop

这对于非常大的m 和i 仍然适用