如何高效实现基于history参数的NumPy数组指定维度堆叠？

给定初始NumPy数据：

import numpy as np
data = np.array([[i for i in range(3)] for _ in range(9)])
print(data)
print(f'data has shape {data.shape}')

输出为：

[[0 1 2]
 [0 1 2]
 [0 1 2]
 [0 1 2]
 [0 1 2]
 [0 1 2]
 [0 1 2]
 [0 1 2]
 [0 1 2]]
data has shape (9, 3)

需要实现以下操作：根据参数history，将连续的history个形状为(3,)的子数组堆叠，形成形状为(N, history, 3)的新数组（N为可完整划分的组数，末尾无法整除的部分直接丢弃）。举例：

• 当history=2时，第一次处理后得到形状(1,2,3)的数组：

history = 2
data = np.array([[[0, 1, 2], [0, 1, 2]]])
print(f'data has now shape {data.shape}')

输出：

data has now shape (1, 2, 3)

• 两次处理后得到形状(2,2,3)的数组，最终处理完成后得到形状(4,2,3)的数组（因9//2=4，最后1个元素被丢弃）。

无需循环，直接利用NumPy的数组切片与重塑操作即可高效完成，代码如下：

import numpy as np

# 原始数据
data = np.array([[i for i in range(3)] for _ in range(9)])
history = 2

# 计算可完整划分的最大行数，截断数据
max_valid_rows = (data.shape[0] // history) * history
processed_data = data[:max_valid_rows].reshape(-1, history, data.shape[1])

# 输出结果
print(processed_data)
print(f'processed data shape: {processed_data.shape}')

运行输出：

[[[0 1 2]
  [0 1 2]]

 [[0 1 2]
  [0 1 2]]

 [[0 1 2]
  [0 1 2]]

 [[0 1 2]
  [0 1 2]]]
processed data shape: (4, 2, 3)

实现说明 #

• data[:max_valid_rows]：直接截断原始数据，保留能被history整除的前N行，丢弃末尾无法组成一组的数据；
• reshape(-1, history, data.shape[1])：-1让NumPy自动计算该维度的长度（即max_valid_rows // history），将一维连续的数组块直接重塑为三维结构，这是NumPy原生优化的操作，执行效率远高于手动循环堆叠。

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