Numpy二维数组按行分组转三维数组的高效索引实现方法

最优实现方案 #

你当前使用的np.apply_along_axis本质是逐行循环迭代，会产生大量额外的数组拷贝，大规模数据集下性能极差，推荐用Numpy步长视图方案实现，全程无数据拷贝，性能可以提升几个数量级。

方案1：Numpy >= 1.20.0 官方滑动窗口API（最稳定）

直接调用官方封装的滑动窗口视图接口，不需要手动计算步长，避免出错：

import numpy as np

# 注意不要用input作为变量名，会覆盖Python内置input函数
arr = np.arange(100) + np.arange(100)[:,None]
window_size = 3

# 生成滑动窗口视图，调整维度顺序匹配你的预期结构
output = np.lib.stride_tricks.sliding_window_view(arr, window_shape=window_size, axis=0).transpose(0,2,1)

# 如果需要独立于原数组的新数组，加.copy()即可：
# output = output.copy()

方案2：低版本Numpy兼容实现

如果你的Numpy版本低于1.20，没有内置滑动窗口接口，可以手动通过as_strided构造视图：

import numpy as np

arr = np.arange(100) + np.arange(100)[:,None]
window_size = 3
rows, cols = arr.shape

# 手动指定数组形状和步长
output = np.lib.stride_tricks.as_strided(
    arr,
    shape=(rows - window_size + 1, window_size, cols),
    strides=(arr.strides[0], arr.strides[0], arr.strides[1])
)

结果验证 #

两种方案生成的结果和你现有代码的输出完全一致，可通过np.allclose(output, 旧实现输出)验证。

内容的提问来源于stack exchange，提问作者SuperCodeBrah