高效的 numpy 索引：获取 M 行的每个块的前 N 行答案

【问题标题】：Efficient numpy indexing: Take first N rows of every block of M rows高效的 numpy 索引：获取 M 行的每个块的前 N 行
【发布时间】：2017-02-03 04:52:10
【问题描述】：

x = np.array([1,2,3,4,5,6,7,8,9,10,11,12,13,14,15])

我想从每个 5 块中获取数组 x 的前 2 行，结果应该是：

x[fancy_indexing] = [1,2, 6,7, 11,12]

使用 for 循环构建这样的索引很容易。

是否有一种单线切片技巧可以成功？此处为简单起见。

【问题讨论】：

标签： python numpy vectorization

【解决方案1】：

方法 #1 这是使用 boolean-indexing 的矢量化单线 -

x[np.mod(np.arange(x.size),M)<N]

方法 #2 如果您追求性能，这里是另一种使用 NumPy strides 的矢量化方法 -

n = x.strides[0]
shp = (x.size//M,N)
out = np.lib.stride_tricks.as_strided(x, shape=shp, strides=(M*n,n)).ravel()

示例运行 -

In [61]: # Inputs
    ...: x = np.array([1,2,3,4,5,6,7,8,9,10,11,12,13,14,15])
    ...: N = 2
    ...: M = 5
    ...: 

In [62]: # Approach 1
    ...: x[np.mod(np.arange(x.size),M)<N]
Out[62]: array([ 1,  2,  6,  7, 11, 12])

In [63]: # Approach 2
    ...: n = x.strides[0]
    ...: shp = (x.size//M,N)
    ...: out=np.lib.stride_tricks.as_strided(x,shape=shp,strides=(M*n,n)).ravel()
    ...: 

In [64]: out
Out[64]: array([ 1,  2,  6,  7, 11, 12])

【讨论】：

【解决方案2】：

import numpy as np

x = np.array(range(1, 16))
y = np.vstack([x[0::5], x[1::5]]).T.ravel()
y
// => array([ 1,  2,  6,  7, 11, 12])

取数组[1, 2, ..., K中每个M行块的第一行N]：

import numpy as np

K = 30
M = 5
N = 2

x = np.array(range(1, K+1))
y = np.vstack([x[i::M] for i in range(N)]).T.ravel()
y
// => array([ 1,  2,  6,  7, 11, 12, 16, 17, 21, 22, 26, 27])

请注意，.T 和 .ravel() 是快速操作：它们不复制任何数据，而只是操纵数组的维度和步幅。

如果您坚持使用精美的索引获取切片：

import numpy as np

K = 30
M = 5
N = 2

x = np.array(range(1, K+1))
fancy_indexing = [i*M+n for i in range(len(x)//M) for n in range(N)]
x[fancy_indexing]
// => array([ 1,  2,  6,  7, 11, 12, 16, 17, 21, 22, 26, 27])

【讨论】：

【解决方案3】：

将数组重新整形为多行，每行五列，然后取（切片）每行的前两列。

>>> x
array([ 1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15])
>>> x.reshape(x.shape[0] / 5, 5)[:,:2]
array([[ 1,  2],
       [ 6,  7],
       [11, 12]])

或者

>>> x.reshape(x.shape[0] / 5, 5)[:,:2].flatten()
array([ 1,  2,  6,  7, 11, 12])
>>>

它仅适用于长度为 5 的倍数的一维数组。

【讨论】：

【解决方案4】：

由于您的措辞是“M 行的每个块的前 N 行”，我首先认为您需要它来处理 2d 数组，所以我将保留我的解决方案。

您可以通过将数组重塑为 3d 来发挥一些作用：

M = 5 # size of blocks
N = 2 # number of columns to cut
x = np.arange(3*4*M).reshape(4,-1) # (4,3*N)-shaped dummy input
x = x.reshape(x.shape[0],-1,M)[:,:,:N+1].reshape(x.shape[0],-1) # (4,3*N)-shaped output

这将根据您的喜好提取每一列。为了将它用于您的 1d 案例，您需要使用 x = x[None,:] 将您的 1d 数组变成 2d 数组。

【讨论】：